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
|