oliver-aizip commited on
Commit
83d0454
·
1 Parent(s): 593a8e7

remove explicit tokenizer device assignment

Browse files
Files changed (1) hide show
  1. utils/models.py +1 -1
utils/models.py CHANGED
@@ -240,7 +240,7 @@ def run_inference(model_name, context, question):
240
  return_tensors="pt",
241
  return_dict=True,
242
  **tokenizer_kwargs,
243
- ).to(device)
244
  with torch.inference_mode():
245
  # Check interrupt before generation
246
  if generation_interrupt.is_set():
 
240
  return_tensors="pt",
241
  return_dict=True,
242
  **tokenizer_kwargs,
243
+ ).to(bitnet_model.device)
244
  with torch.inference_mode():
245
  # Check interrupt before generation
246
  if generation_interrupt.is_set():