fix v dtype in ema black list (#8466)

PaddlePaddle · Jul 25, 2023 · 47d7ee5 · 47d7ee5
1 parent 9ed7ad2
commit 47d7ee5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ppdet/optimizer/ema.py b/ppdet/optimizer/ema.py
@@ -69,7 +69,7 @@ def __init__(self,
         self.state_dict = dict()
         for k, v in model.state_dict().items():
             if k in self.ema_black_list:
-                self.state_dict[k] = v.astype('float32')
+                self.state_dict[k] = v
             else:
                 self.state_dict[k] = paddle.zeros_like(v, dtype='float32')
 
@@ -127,7 +127,7 @@ def apply(self):
         for k, v in self.state_dict.items():
             if k in self.ema_black_list:
                 v.stop_gradient = True
-                state_dict[k] = v.astype(model_dict[k].dtype)
+                state_dict[k] = v
             else:
                 if self.ema_decay_type != 'exponential':
                     v = v / (1 - self._decay**self.step)