Update megatron/training_log.py

argonne-lcf · Sep 13, 2024 · 7f71572 · 7f71572
1 parent 2113dbc
commit 7f71572
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/training_log.py b/megatron/training_log.py
@@ -397,7 +397,7 @@ def training_log(
                     opt_stats, group=mpu.get_pipeline_model_parallel_group()
                 )
                 # opt_stats_2 = opt_stats_2.clone().detach()
-                opt_stats_2 = accelerator.get_accelerator().FloatTensor(opt_stats_2)
+                opt_stats_2 = accelerator.FloatTensor(opt_stats_2)
                 torch.distributed.all_reduce(
                     opt_stats_2,
                     op=torch.distributed.ReduceOp.MAX,