Update app.py

app.py

@@ -13,7 +13,7 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
-model = model.to('cuda:0')
+#model = model.to('cuda:0')
 
 
 def make_inference(product_name, product_description):
@@ -21,6 +21,8 @@ def make_inference(product_name, product_description):
         f"### Product and Description:\n{product_name}: {product_description}\n\n### Ad:",
         return_tensors="pt",
     )
+    # Move the input tensors to the GPU
+    batch = {key: value.to('cuda:0') for key, value in batch.items()}
 
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=50)