Update app.py

app.py

@@ -181,7 +181,7 @@ def __(torch):
                 item = torch.mean(v, dim=0)
         attns.append(item + int(skip_connection) * torch.eye(item.shape[0]))
         roll = torch.prod(torch.stack(attns), dim=0)
-        return roll.numpy()
+        return roll
 
     return (rollout,)