return grad in post_backward_hook()

facebookresearch · May 8, 2024 · ad40f24 · ad40f24
1 parent 901fb86
commit ad40f24
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
@@ -411,8 +411,8 @@ def _hook(
             self.fp32_grads[param_index] = grad.to(torch.float32)
         else:
             self.fp32_grads[param_index].add_(grad.data)
-
         #logger.info(f"CHRISLOG: after post-backward hook, self.fp32_grads[param_index] is None: {self.fp32_grads[param_index] is None}")
+        return grad
 
     def _unflatten_params_as_views(self) -> None:
         """Unlike ``_unflatten_params``, this function unflatten into views and keep