post rebase

facebookresearch · Dec 4, 2023 · 9791920 · 9791920
1 parent dba4758
commit 9791920
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1777,9 +1777,9 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # case grads should be all-reduced here.
                 assert self.world_size == 1
                 if getattr(param, "main_grad", None) is not None:
-                    self._post_reduction_hook(param, param.main_grad.data)
+                    self._post_reduction_hook(param, param.main_grad)
                 else:
-                    self._post_reduction_hook(param, param.grad.data)
+                    self._post_reduction_hook(param, param.grad)
 
             # After _post_backward_hook returns, orig_grad_data will eventually
             # go out of scope, at which point it could otherwise be freed for