Update app.py

app.py

@@ -23,13 +23,13 @@ model_layer_names = ["0", "1", "2", "3"]
 
 def get_layer(layer_name):
     print(layer_name)
-    if layer_name == "0":
+    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
@@ -92,7 +92,7 @@ demo = gr.Interface(
         gr.Image(width=256,height=256,label="Input image"),
         gr.Number(value=3, maximum=10, minimum=1,step=1.0, precision=0,label="Number of classes to display"),
         gr.Checkbox(True, label="Show GradCAM Image"),
-        gr.Dropdown(model_layer_names, value="3", label="Which layer for Gradcam"),
+        gr.Dropdown(model_layer_names, value=3, label="Which layer for Gradcam"),
         gr.Slider(0, 1, value=0.5,label="Overall opacity of the overlay"),
     ],
     outputs = [
@@ -102,8 +102,8 @@ 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", "1", True, 10, -1],
-                ["dog.jpg", "1", False,  4, -1]]
+    examples = [["cat.jpg", 1, True, 10, -1],
+                ["dog.jpg", 1, False,  4, -1]]
 )
 
 if __name__ == "__main__":