guard main_grad None

facebookresearch · Oct 8, 2023 · 0be2e5e · 0be2e5e
1 parent bb516bf
commit 0be2e5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fairscale/nn/misc/flatten_params_wrapper.py b/fairscale/nn/misc/flatten_params_wrapper.py
@@ -370,7 +370,7 @@ def _unflatten_params_as_views(self) -> None:
         """
         assert self.is_flattened
         for p in self.flat_params:
-            if not hasattr(p, 'main_grad') or p.main_grad.shape != p.shape:
+            if getattr(p, 'main_grad', None) is None or p.main_grad.shape != p.shape:
                 p.main_grad = torch.zeros_like(p, dtype=torch.float32)
         ps, ps_main_grad = self.get_param_views()
         param_views = []