🔇🐛Fixed fit

carefree0910 · Oct 14, 2024 · 7c8f874 · 7c8f874
1 parent a3d4865
commit 7c8f874
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/core/learn/trainer.py b/core/learn/trainer.py
@@ -350,10 +350,13 @@ def fit(
             else:
                 final_score = self.final_results.final_score
             self.save_checkpoint(final_score)
-        console.debug("finalizing training")
+        if self.is_local_rank_0:
+            console.debug("finalizing training")
         for callback in self.callbacks:
             callback.finalize(self)
-        console.debug("finished training")
+        if self.is_local_rank_0:
+            console.debug("finished training")
+        self.accelerator.wait_for_everyone()
         return self
 
     ## checkpointing