"#35049 fix path for num_shard 1"

huggingface · Dec 5, 2024 · 8208593 · 8208593
1 parent 3129967
commit 8208593
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/transformers/models/mllama/convert_mllama_weights_to_hf.py b/src/transformers/models/mllama/convert_mllama_weights_to_hf.py
@@ -338,7 +338,11 @@ def write_model(
 
     print(f"Fetching all parameters from the checkpoint at {input_base_path}...")
     if num_shards == 1:
-        loaded = [torch.load(os.path.join(input_base_path, "consolidated.pth"), map_location="cpu", mmap=True)]
+        if os.path.exists(os.path.join(input_base_path, "consolidated.00.pth")):
+            path = os.path.join(input_base_path, "consolidated.00.pth")
+        else:
+            path = os.path.join(input_base_path, "consolidated.pth")
+        loaded = [torch.load(path, map_location="cpu", mmap=True)]
     else:
         loaded = [
             torch.load(os.path.join(input_base_path, f"consolidated.{i:02d}.pth"), map_location="cpu", mmap=True)