From 15b3b043f97408e9c9cb3d3a988ebdc04417265b Mon Sep 17 00:00:00 2001 From: Shriya Palsamudram <69161273+ShriyaPalsamudram@users.noreply.github.com> Date: Mon, 15 Apr 2024 20:06:45 -0400 Subject: [PATCH] Update tranpose call to match TE API change (#8918) Signed-off-by: Shriya Palsamudram --- nemo/core/optim/distributed_adam.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/core/optim/distributed_adam.py b/nemo/core/optim/distributed_adam.py index 32bd7e6c1154..f58fb7352c38 100644 --- 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()