Update

[ghstack-poisoned]
pytorch · Feb 26, 2025 · 206df62 · 206df62
1 parent 6368448
commit 206df62
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/torchtitan/components/checkpoint.py b/torchtitan/components/checkpoint.py
@@ -314,12 +314,13 @@ def __init__(
         )
 
     def __del__(self):
-        if self.enable_checkpoint and self.mp and self.mp.is_alive():
-            self.mp_queue_send.put(Terminate())
-            self.mp.join()
-        if self.purge_thread and self.purge_thread.is_alive():
-            self.purge_queue.put(Terminate())
-            self.purge_thread.join()
+        if self.enable_checkpoint
+            if self.mp and self.mp.is_alive():
+                self.mp_queue_send.put(Terminate())
+                self.mp.join()
+            if self.purge_thread and self.purge_thread.is_alive():
+                self.purge_queue.put(Terminate())
+                self.purge_thread.join()
 
     @torch.no_grad()
     def save(self, curr_step: int, force: bool = False) -> None: