fixed typo in fp8 weight tensor name

Signed-off-by: Alp Dener <[email protected]>
NVIDIA · Mar 11, 2024 · 9d8a7f5 · 9d8a7f5
1 parent 7627bfa
commit 9d8a7f5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
@@ -167,7 +167,7 @@ def forward(
                     or (is_fp8_activation_recompute_enabled()
                         and not in_fp8_activation_recompute_phase())):
                     # Gather Fp8 transposed-weight buffers if needed
-                    weight_t_fp8_shape = tuple(reversed(weight_t_fp8._data))
+                    weight_t_fp8_shape = tuple(reversed(weight_t_fp8._data.shape))
                     if (fsdp_group is not None
                         and weight_t_fp8._data.shape != weight_t_fp8_shape):
                         _fsdp_gather_tensors(fsdp_group,