Update app.py

app.py

@@ -22,13 +22,14 @@ classes = ('plane', 'car', 'bird', 'cat', 'deer',
 model_layer_names = ["0", "1", "2", "3"]
 
 def get_layer(layer_name):
-    if layer_name == 0:
+    print(layer_name)
+    if layer_name == "0":
         return [model.prep[-1]]
-    elif layer_name == 1:
+    elif layer_name == "1":
         return [model.layer1[-1]]
-    elif layer_name == 2:
+    elif layer_name == "2":
         return [model.layer2[-1]]
-    elif layer_name == 3:
+    elif layer_name == "3":
         return [model.layer3[-1]]
     else:
         return None
@@ -101,7 +102,7 @@ demo = gr.Interface(
     title = "CIFAR 10 trained on ResNet model in pytorch lightning with Gradcam",
     description = " A simple gradio inference to infer on resnet18 model",
     examples = [["cat.jpg", 10, True, "1", -1],
-                ["dog.jpg", 4,  False, "2", -1]]
+                ["dog.jpg", 4,  False, "1", -1]]
 )
 
 if __name__ == "__main__":