danilohssantana commited on
Commit
1af9e28
·
1 Parent(s): bcf356c

fixing device problem

Browse files
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -129,7 +129,7 @@ def predict(data: PredictRequest):
129
  videos=video_inputs,
130
  padding=True,
131
  return_tensors="pt",
132
- ).to("cuda:0")
133
 
134
  # Generate the output
135
  generated_ids = model.generate(**inputs, max_new_tokens=2056)
 
129
  videos=video_inputs,
130
  padding=True,
131
  return_tensors="pt",
132
+ ).to(model.device)
133
 
134
  # Generate the output
135
  generated_ids = model.generate(**inputs, max_new_tokens=2056)