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

Change to only print temp directory once on rank0.

Pull Request resolved: #89964
Approved by: https://github.com/XilunWu

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: