fix weight_map

huggingface · Jul 18, 2024 · 8ef00e0 · 8ef00e0
1 parent 8c16267
commit 8ef00e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/fx/parallelization/api.py b/optimum/fx/parallelization/api.py
@@ -118,7 +118,7 @@ def parallelize_model(
             if os.path.isfile(index_path):
                 with open(index_path) as f:
                     index_dict = json.load(f)
-                parallel_ctx.weight_map = index_dict["weight_map"]
+                parallel_ctx.weight_map = {k : os.path.join(hf_folder, v) for k, v in index_dict["weight_map"].items()}
             weight_files = glob.glob(os.path.join(hf_folder, "*.safetensors" if use_safetensors else "*.bin"))
             if not use_safetensors:
                 weight_map = parallel_ctx.weight_map if parallel_ctx.weight_map else {}