Update donut_inference.py

donut_inference.py

@@ -10,10 +10,10 @@ load_dotenv()
 # imgae = image.resize((1864, 1440))
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-device = "cpu"
+# device = "cpu"
 # Load the processor from the local directory
 processor = DonutProcessor.from_pretrained("Henge-navuuu/donut-base-finetuned-forms-v1")
-
+processor.to(device)
 # Load the model from the local directory
 model = VisionEncoderDecoderModel.from_pretrained("Henge-navuuu/donut-base-finetuned-forms-v1")
 model.to(device)
@@ -25,7 +25,7 @@ def inference(image):
     decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt")["input_ids"]
 
     # device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
+    # model.to(device)
 
     outputs = model.generate(pixel_values.to(device),
                                 decoder_input_ids=decoder_input_ids.to(device),