Update app.py

app.py

@@ -4,10 +4,17 @@ import gradio as gr
 from torchvision.transforms import Compose, Normalize, ToTensor, Resize, CenterCrop
 
 
-inference = torch.load('fine_tune_resnet.pth')
+inference = torch.load('fine_tune_resnet.pth', map_location=torch.device('cpu'))
 inference.eval()
 
 def classifier(image):
+    test_transform = Compose([
+            Resize(256),
+            CenterCrop(224),
+            ToTensor(),
+            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
     if torch.cuda.is_available():
         inference.cpu()