Update app.py

app.py

@@ -19,13 +19,11 @@ inv_normalize = transforms.Normalize(
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')
 
-model_layer_names = ["0", "1", "2", "3"]
+model_layer_names = [0, 1, 2, 3]
 
 def get_layer(layer_name):
     print(layer_name)
-    if layer_name == 0:
-        return [model.prep[-1]]
-    elif layer_name == 1:
+    if layer_name == 1:
         return [model.layer1[-1]]
     elif layer_name == 2:
         return [model.layer2[-1]]