oluinioluwa814 commited on
Commit
9af31ae
·
verified ·
1 Parent(s): 7e05315

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +1 -1
model.py CHANGED
@@ -7,7 +7,7 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
7
  model = AutoModelForCausalLM.from_pretrained(
8
  MODEL_NAME,
9
  device_map="auto",
10
- torch_dtype=torch.float16
11
  )
12
 
13
  def generate_response(prompt, max_tokens=512):
 
7
  model = AutoModelForCausalLM.from_pretrained(
8
  MODEL_NAME,
9
  device_map="auto",
10
+ dtype=torch.float16
11
  )
12
 
13
  def generate_response(prompt, max_tokens=512):