Remove unnecessary torch stack

Author: Ian Schweer

pytensor/link/pytorch/dispatch/scalar.py

@@ -68,7 +68,10 @@ def scalar_loop(steps, *start_and_constants):
                 *carry, done = update(*carry, *constants)
                 if done:
                     break
-            return torch.stack(carry), torch.tensor([done])
+            if len(node.outputs) == 2:
+                return carry[0], done
+            else:
+                return carry, done
     else:
 
         def scalar_loop(steps, *start_and_constants):
@@ -78,6 +81,9 @@ def scalar_loop(steps, *start_and_constants):
             )
             for _ in range(steps):
                 carry = update(*carry, *constants)
-            return torch.stack(carry)
+            if len(node.outputs) == 1:
+                return carry[0]
+            else:
+                return carry
 
     return torch.compiler.disable(scalar_loop)