Update app.py

app.py

@@ -1,3 +1,4 @@
+import torch
 import spaces
 import gradio as gr
 from transformers import pipeline
@@ -5,7 +6,7 @@ from PIL import Image
 from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 
 # Load model and processor
-model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-screen2words-large", torch_dtype=torch.bfloat16).to("cuda")
+model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-screen2words-large", dtype=torch.bfloat16).to("cuda")
 processor = Pix2StructProcessor.from_pretrained("google/pix2struct-screen2words-large")
 
 # Define the function