fixed loading of state dict

app.py

@@ -6,7 +6,8 @@ from torchvision.transforms import Resize, ToTensor
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
 model = torch.nn.Module()
-model.load_state_dict('model.pt', map_location=device)
+
+model.load_state_dict(torch.load('model.pt', map_location=device))
 
 resize = Resize((224))
 to_tensor = ToTensor()