Address pr comments

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/elemwise.py

@@ -21,13 +21,11 @@ def elemwise_fn(*inputs):
 
         def elemwise_fn(*inputs):
             Elemwise._check_runtime_broadcast(node, inputs)
-            shaped_inputs = torch.broadcast_tensors(*inputs)
+            broadcast_inputs = torch.broadcast_tensors(*inputs)
             ufunc = base_fn
-            for _ in range(shaped_inputs[0].dim()):
+            for _ in range(broadcast_inputs[0].dim()):
                 ufunc = torch.vmap(ufunc)
-            # @todo: This will fail for anything that calls
-            # `.item()`
-            return ufunc(*shaped_inputs)
+            return ufunc(*broadcast_inputs)
 
     return elemwise_fn