zqh11 commited on
Commit
9a4db57
1 Parent(s): fd806a6

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -144,7 +144,7 @@ from model import IrisClassifier as Classifier
144
  def main():
145
  # Model training and evaluation
146
  """
147
- inputs = tokenizer(input_text, return_tensors="pt").cuda()
148
  outputs = model.generate(**inputs, max_new_tokens=140)
149
  print(tokenizer.decode(outputs[0]))
150
  ```
 
144
  def main():
145
  # Model training and evaluation
146
  """
147
+ inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
148
  outputs = model.generate(**inputs, max_new_tokens=140)
149
  print(tokenizer.decode(outputs[0]))
150
  ```