david167 commited on
Commit
c86959d
·
1 Parent(s): 01a04bc

Fix remaining device_map auto in gradio_app.py

Browse files
Files changed (1) hide show
  1. gradio_app.py +1 -1
gradio_app.py CHANGED
@@ -56,7 +56,7 @@ class ModelManager:
56
  self.model = AutoModelForCausalLM.from_pretrained(
57
  base_model_name,
58
  torch_dtype=torch.float16 if self.device == "cuda:0" else torch.float32,
59
- device_map="auto" if self.device == "cuda:0" else None,
60
  trust_remote_code=True,
61
  low_cpu_mem_usage=True,
62
  token=hf_token
 
56
  self.model = AutoModelForCausalLM.from_pretrained(
57
  base_model_name,
58
  torch_dtype=torch.float16 if self.device == "cuda:0" else torch.float32,
59
+ device_map={"": 0} # Force all parameters to GPU 0,
60
  trust_remote_code=True,
61
  low_cpu_mem_usage=True,
62
  token=hf_token