wrong device issue #1196 (#1197)

ModelCloud · Feb 1, 2025 · f095cb0 · f095cb0
1 parent 9e4129c
commit f095cb0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/gptqmodel/models/loader.py b/gptqmodel/models/loader.py
@@ -98,6 +98,8 @@ def get_model_local_path(pretrained_model_id_or_path, **kwargs):
     if is_local:
         return pretrained_model_id_or_path
     else:
+        # hf_transfer does not accept max_memory arg
+        kwargs.pop('max_memory', None)
         return snapshot_download(pretrained_model_id_or_path, **kwargs)
 
 def get_tokenizer(model_id_or_path, config, trust_remote_code: bool = False):

diff --git a/requirements.txt b/requirements.txt
@@ -10,3 +10,4 @@ device-smi==0.3.3
 sentencepiece>=0.2.0
 protobuf>=5.29.1
 pillow>=10.4.0
+hf_transfer>=0.1.9