Fix the contiguous_param_buffer bug about bprop overlap and redundant…

… copy after all-gather.

nemo/core/optim/distributed_adam.py

@@ -371,11 +371,13 @@ def init_param_buffer(self) -> None:
                 param._data = buffer_view.view(param.size())
             else:
                 # Preserve memory format for param here, i.e. NHWC tensors
-                param.data.set_(
-                    source=buffer_view,
-                    storage_offset=0,
-                    size=param.size(),
-                    stride=param.stride(),
+                # `param.data.set_()` failed to change storage.
+                # `param.set_()` invalidates bprop hook.
+                param.data = torch.as_strided(
+                    buffer_view,
+                    param.size(),
+                    param.stride(),
+                    storage_offset=buffer_view.storage_offset(),
                 )
 
     def try_grad_sync(self, params: Iterable[torch.nn.Parameter]) -> None: