sharded files fix

neuralmagic · Aug 20, 2024 · bc0caac · bc0caac
1 parent 3546d74
commit bc0caac
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/compressed_tensors/utils/safetensors_load.py b/src/compressed_tensors/utils/safetensors_load.py
@@ -47,6 +47,10 @@ def get_safetensors_folder(
     model will be searched for in the default TRANSFORMERS_CACHE
     :return: local folder containing model data
     """
+    if isinstance(pretrained_model_name_or_path, list):
+        # assume sharded files, referencing first file is sufficient
+        pretrained_model_name_or_path = pretrained_model_name_or_path[0]
+
     if os.path.exists(pretrained_model_name_or_path):
         # argument is a path to a local folder
         return os.path.abspath(pretrained_model_name_or_path)