Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import torchvision.transforms as transforms
 import torch
 
-def transform_image(image):
+def transform_image(device, image):
     my_transforms = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize(
                                             [0.485, 0.456, 0.406],
@@ -19,7 +19,7 @@ def predict(img):
     model.eval()
     
     image = transforms.Resize((480,640))(Image.fromarray(img))
-    tensor = transform_image(image=image)
+    tensor = transform_image(device, image)
     
     model.to(device)
     with torch.no_grad():