Update hf_auto_model_for_causal_lm.py

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Feb 27, 2025 · e8bd0fa · e8bd0fa
1 parent c69fdf6
commit e8bd0fa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/llm/gpt/model/hf_auto_model_for_causal_lm.py b/nemo/collections/llm/gpt/model/hf_auto_model_for_causal_lm.py
@@ -290,7 +290,7 @@ def on_before_optimizer_step(self, optimizer) -> None:
             mean_loss = torch.tensor([mean_loss], device=self.device).detach()
             dist.all_reduce(mean_loss, group=group, op=dist.ReduceOp.AVG)
             mean_loss = mean_loss.item()
-            tps = torch.tensor([tps], device=self.device, dtype=torch.int64)
+            tps = torch.tensor([tps], device=self.device, dtype=torch.int64).detach()
             dist.all_reduce(tps, group=group, op=dist.ReduceOp.SUM)
             tps = tps.item()