fix(nyz): fix rmsprop bug in torch 1.13.1

opendilab · Dec 13, 2024 · 580ea65 · 580ea65
1 parent 317e775
commit 580ea65
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ding/torch_utils/optimizer_helper.py b/ding/torch_utils/optimizer_helper.py
@@ -495,7 +495,7 @@ def _state_init(self, p, momentum, centered):
             # wait torch upgrad to 1.4, 1.3.1 didn't support memory format state['step'] = 0
         else:
             state['step'] = torch.zeros((1,), dtype=torch.float, device=p.device) \
-                if self.defaults['capturable'] else torch.tensor(0.)
+                if ('capturable' in self.defaults and self.defaults['capturable']) else torch.tensor(0.)
         state['thre_square_avg'] = torch.zeros_like(p.data, device=p.data.device)
         state['square_avg'] = torch.zeros_like(p.data, device=p.data.device)
         if momentum: