Update app.py
Browse files
app.py
CHANGED
|
@@ -21,7 +21,8 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
| 21 |
model = AutoModelForCausalLM.from_pretrained(model_id,
|
| 22 |
# load_in_8bit=True,
|
| 23 |
# quantization_config=nf4_config
|
| 24 |
-
torch_dtype = torch.bfloat16
|
|
|
|
| 25 |
)
|
| 26 |
|
| 27 |
@spaces.GPU
|
|
|
|
| 21 |
model = AutoModelForCausalLM.from_pretrained(model_id,
|
| 22 |
# load_in_8bit=True,
|
| 23 |
# quantization_config=nf4_config
|
| 24 |
+
torch_dtype = torch.bfloat16,
|
| 25 |
+
device_map="auto"
|
| 26 |
)
|
| 27 |
|
| 28 |
@spaces.GPU
|