Update app.py

app.py

@@ -72,7 +72,7 @@ def visualize_heatmap(model, image):
     model.eval()
     x = image#.unsqueeze(0) # remove batch dimension
     features = model.convnet(x) # feature heatmap learnt by model
-    heatmap = torch.mean(features, dim=1).squeeze().detach().numpy() # normalize heatmap to ndarray
+    heatmap = torch.mean(features, dim=1).detach().numpy() #.squeeze() normalize heatmap to ndarray
     plt.imshow(heatmap, cmap="hot") # display heatmap as plot
     plt.axis("off")