dist adam transpose fix

Signed-off-by: dimapihtar <[email protected]>
NVIDIA · May 17, 2024 · ff0eb2f · ff0eb2f
1 parent 6a5187b
commit ff0eb2f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/core/optim/distributed_adam.py b/nemo/core/optim/distributed_adam.py
@@ -487,7 +487,7 @@ def _param_copy_fragments(self, fragments: Iterable[DistributedFusedAdam.Paramet
         for param in params:
             if is_float8tensor(param):
                 param._reset_caches()
-                param.transpose(update_cache=True)
+                param.transpose_2d(cache=True)
                 param._lazy_transpose_cache = True
 
     @torch.no_grad()