Update app.py

app.py

@@ -218,13 +218,13 @@ def __(dropdown_method, focus_square, rollout, switch, torch):
             a = act[layer_key][0, :, ::-1 , :]
         elif dropdown_method.value == "Attention rollout (MIN)":
             a = rollout(act, skip_connection=switch.value, parse="min")
-            a = torch.stack([a for _ in range(32)])
+            a = torch.stack([a for _ in range(32)]).numpy()
         elif dropdown_method.value == "Attention rollout (MAX)":
             a = rollout(act, skip_connection=switch.value, parse="max")
-            a = torch.stack([a for _ in range(32)])
+            a = torch.stack([a for _ in range(32)]).numpy()
         elif dropdown_method.value == "Attention rollout (MEAN)":
             a = rollout(act, skip_connection=switch.value, parse="mean")
-            a = torch.stack([a for _ in range(32)])
+            a = torch.stack([a for _ in range(32)]).numpy()
         return a
     focus_square_ind = 8 * (int(focus_square.value[1]) - 1) + ord(focus_square.value[0]) - ord("a")