Spaces:
Sleeping
Sleeping
Fix remaining device_map auto in gradio_app.py
Browse files- 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="
|
| 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
|