Update app.py

app.py

@@ -12,26 +12,20 @@ import re
 import matplotlib.pyplot as plt
 from io import BytesIO
 
-# inv_normalize = transforms.Normalize(
-#     mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
-#     std=[1/0.23, 1/0.23, 1/0.23]
-# )
-
-inv_normalize = transforms.Normalize(
-    mean=[0.49139968, 0.48215827 ,0.44653124],
-    std=[0.24703233, 0.24348505, 0.26158768]
-)
+transform = transforms.Compose([
+      transforms.ToTensor(),
+      transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768))])
 
 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
 
 model = LITResNet(classes)
-model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
+model.load_state_dict(torch.load("model.pth")["state_dict"])
+model.eval()
 modellayers = list(dict(model.named_modules()))
 
 def inference(input_img, num_gradcam_images=1, target_layer_number=-1, transparency=0.5, show_misclassified=False, num_top_classes=3, num_misclassified_images=3):
     input_img = np.array(Image.fromarray(np.array(input_img)).resize((32, 32)))
     org_img = input_img
-    transform = transforms.ToTensor()
     input_img = transform(input_img).unsqueeze(0)
     outputs = model(input_img)
     softmax = torch.nn.Softmax(dim=0)