tabedini commited on
Commit
0096c75
1 Parent(s): 4b4284a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -60,7 +60,7 @@ if not torch.cuda.is_available():
60
 
61
  if torch.cuda.is_available():
62
  model_id = "PartAI/Dorna-Llama3-8B-Instruct"
63
- model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
64
  tokenizer = AutoTokenizer.from_pretrained(model_id)
65
 
66
  generation_speed = 0
 
60
 
61
  if torch.cuda.is_available():
62
  model_id = "PartAI/Dorna-Llama3-8B-Instruct"
63
+ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16)
64
  tokenizer = AutoTokenizer.from_pretrained(model_id)
65
 
66
  generation_speed = 0