bugfix: fix #971

GaiZhenbiao · Dec 25, 2023 · 6e50743 · 6e50743
1 parent 3abd6e1
commit 6e50743
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/modules/models/Qwen.py b/modules/models/Qwen.py
@@ -21,7 +21,7 @@ def __init__(self, model_name, user_name="") -> None:
  except KeyError:
  model_source = model_name
  self.tokenizer = AutoTokenizer.from_pretrained(model_source, trust_remote_code=True, resume_download=True)
- self.model = AutoModelForCausalLM.from_pretrained(model_source, device_map="auto", trust_remote_code=True, resume_download=True).eval()
+ self.model = AutoModelForCausalLM.from_pretrained(model_source, device_map="cuda", trust_remote_code=True, resume_download=True).eval()
 
  def generation_config(self):
  return GenerationConfig.from_dict({

diff --git a/requirements.txt b/requirements.txt
@@ -33,3 +33,4 @@ pydantic==2.5.2
 google-search-results
 anthropic==0.3.11
 Pillow>=10.1.0
+protobuf==3.20.3