Wrong flags

pymc-devs · Jan 26, 2025 · dbc95e4 · dbc95e4
1 parent 7b4c426
commit dbc95e4
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/pytensor/link/pytorch/linker.py b/pytensor/link/pytorch/linker.py
@@ -37,9 +37,8 @@ def conversion_func_register(*args, **kwargs):
     def jit_compile(self, fn):
         import torch
 
-        # two flags that tend to help our graphs
-        torch._dynamo.config.capture_func_transforms = True
-        torch._dynamo.config.capture_scalar_outputs = True
+        # flag that tend to help our graphs
+        torch._dynamo.config.capture_dynamic_output_shape_ops = True
 
         from pytensor.link.pytorch.dispatch import pytorch_typify