bug fix

app.py

@@ -119,11 +119,14 @@ def get_activations(model, image: list, model_name: str,
     layer_outputs = {}
     for i in range(len(model.model.features)):
         image = model.model.features[i](image)
-        layer_outputs[i] = image
-        print(i, layer_outputs[i].shape)
+        if i in activation_indices[model_name]:
+            layer_outputs[i] = image
+            print(i, layer_outputs[i].shape)
+        if i == max(activation_indices[model_name]):
+            break
     output = model.model(image).detach().cpu().numpy()
-    output_1 = activation_indices[model_name].detach().cpu().numpy()
-    output_2 = activation_indices[model_name].detach().cpu().numpy()
+    output_1 = activation_indices[model_name][0].detach().cpu().numpy()
+    output_2 = activation_indices[model_name][1].detach().cpu().numpy()
     
     # get activations
     # output_1 = outputs[1]