torchrun for DDP

agi-brain · Sep 20, 2024 · 046666c · 046666c
1 parent df5cf7c
commit 046666c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/xuance/torch/agents/base/agent.py b/xuance/torch/agents/base/agent.py
@@ -242,6 +242,7 @@ def finish(self):
         if self.distributed_training:
             if dist.get_rank() == 0:
                 if os.path.exists(self.learner.snapshot_path):
-                    os.remove(os.path.join(self.learner.snapshot_path, "snapshot.pt"))
+                    if os.path.exists(os.path.join(self.learner.snapshot_path, "snapshot.pt")):
+                        os.remove(os.path.join(self.learner.snapshot_path, "snapshot.pt"))
                     os.removedirs(self.learner.snapshot_path)
             destroy_process_group()