Update trainer.py

zkyseu · Dec 19, 2022 · 1440efc · 1440efc
1 parent 76b9eec
commit 1440efc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pplanedet/engine/trainer.py b/pplanedet/engine/trainer.py
@@ -170,7 +170,8 @@ def __init__(self, cfg):
             scale_loss = amp_cfg.pop('scale_loss')
             self.scaler = paddle.amp.GradScaler(init_loss_scaling=scale_loss)
             amp_cfg['models'] = self.model
-            self.model = paddle.amp.decorate(**amp_cfg)  # decorate for level O2
+            if amp_cfg['level'] == 'O2':
+                self.model = paddle.amp.decorate(**amp_cfg)  # decorate for level O2
 
         # ZeRO
         self.sharding_strategies = cfg.get('sharding', False)