exist ok diag subf

pytorch · Feb 6, 2025 · 476c5a6 · 476c5a6
1 parent 763f60e
commit 476c5a6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/ibm_rescaling/rescaling_demo.py b/examples/ibm_rescaling/rescaling_demo.py
@@ -138,7 +138,7 @@
     vals = dist.tensor.DTensor.from_local(vals, mesh, placement).full_tensor()
 
     # Diag save
-    os.makedirs(os.path.join(args.ckpt_path, "diag"))
+    os.makedirs(os.path.join(args.ckpt_path, "diag"), exist_ok=True)
     torch.save(data.state_dict(), os.path.join(args.ckpt_path, "diag", f"loader_state_{rank}.pth"))
     time.sleep(10)