small fix

vue1999 · Dec 24, 2024 · eba347a · eba347a
1 parent be846c0
commit eba347a
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/mace/tools/train.py b/mace/tools/train.py
@@ -458,10 +458,9 @@ def closure():
         torch.distributed.barrier()
         return total_loss
 
-    if torch.distributed.get_rank() == 0:
-        loss = optimizer.step(closure)
-        for param in model.parameters():
-            torch.distributed.broadcast(param.data, src=0)
+    loss = optimizer.step(closure)
+    for param in model.parameters():
+        torch.distributed.broadcast(param.data, src=0)
 
     if ema is not None:
         ema.update()