Address pr comments

pymc-devs · Nov 6, 2024 · 97d6bdc · 97d6bdc
1 parent a095076
commit 97d6bdc
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/pytensor/link/pytorch/dispatch/elemwise.py b/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