[Checkpoint] Minor update to checkpoint utils (#89964)

wz337 · pytorchmergebot · commit ad1585b4a45c · 2022-12-01T01:55:53.000Z
Change to only print temp directory once on rank0. Pull Request resolved: #89964 Approved by: https://github.com/XilunWu
diff --git a/torch/testing/_internal/distributed/checkpoint_utils.py b/torch/testing/_internal/distributed/checkpoint_utils.py
@@ -21,13 +21,15 @@ def wrapper(self, *args: Tuple[object], **kwargs: Dict[str, Any]) -> None:
         # Only create temp_dir when rank is 0
         if dist.get_rank() == 0:
             temp_dir = tempfile.mkdtemp()
+            print(f"Using temp directory: {self.temp_dir }")
         else:
             temp_dir = ""
         object_list = [temp_dir]
+
         # Broadcast temp_dir to all the other ranks
         dist.broadcast_object_list(object_list)
         self.temp_dir = object_list[0]
-        print(f"Using temp directory: {self.temp_dir }")
+
         try:
             func(self)
         finally:
