add .to(device) to attention_mask

app.py

@@ -35,7 +35,7 @@ def model_call(question, document, answer):
     NEW_FORMAT = PROMPT.format(question=question, document=document, answer=answer)
     inputs = tokenizer(NEW_FORMAT, return_tensors="pt")
     input_ids = inputs.input_ids.to(device)
-    attention_mask = inputs.attention_mask
+    attention_mask = inputs.attention_mask.to(device)
     generate_kwargs = dict(
         input_ids=input_ids,
         do_sample=True,