Update app.py

app.py

@@ -20,6 +20,9 @@ def make_inference(product_name, product_description):
         return_tensors="pt",
     )
 
+    # Move batch to the same device as the model
+    batch = {k: v.to(model.device) for k, v in batch.items()}
+
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=50)