fix tensor data update for gemini loss caluculation (hpcaitech#5442)

flybird11111 · Mar 18, 2024 · cced72d · cced72d
1 parent 9052280
commit cced72d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/applications/Colossal-LLaMA-2/train.py b/applications/Colossal-LLaMA-2/train.py
@@ -56,6 +56,7 @@ def format_numel_str(numel: int) -> str:
 
 def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
     dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
+    tensor = tensor.data
     tensor.div_(dist.get_world_size())
     return tensor