Fixes format utils executable (#123407)

LucasLLC · pytorchbot · commit 1ab823a68384 · 2024-04-05T22:36:16.000Z
Fixes an issue with the format utils executable, which was causing it to run as a no-op. :( Pull Request resolved: #123407 Approved by: https://github.com/wz337, https://github.com/fegin (cherry picked from commit 18c9d46)
diff --git a/torch/distributed/checkpoint/format_utils.py b/torch/distributed/checkpoint/format_utils.py
@@ -240,7 +240,6 @@ def dcp_to_torch_save(
         To avoid OOM, it's recommended to only run this function on a single rank.
     """
     sd: STATE_DICT_TYPE = {}
-
     _load_state_dict(
         sd,
         storage_reader=FileSystemReader(dcp_checkpoint_dir),
@@ -298,13 +297,15 @@ class FormatMode(Enum):
     checkpoint_missing_warning = (
         f"No checkpoint found at {args.src}. Skipping conversion."
     )
-    if args.mode == FormatMode.TORCH_TO_DCP:
+    if args.mode == FormatMode.TORCH_TO_DCP.value:
         if os.path.isfile(args.src):
             torch_save_to_dcp(args.src, args.dst)
         else:
             print(checkpoint_missing_warning)
-    elif args.mode == FormatMode.DCP_TO_TORCH:
+    elif args.mode == FormatMode.DCP_TO_TORCH.value:
         if os.path.isdir(args.src):
             dcp_to_torch_save(args.src, args.dst)
         else:
             print(checkpoint_missing_warning)
+    else:
+        raise ValueError(f"Unknown conversion mode: {args.mode}")
