Add compiler flags to help hint more to torch

pymc-devs · Jan 17, 2025 · c5f26fd · c5f26fd
1 parent 70634f1
commit c5f26fd
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/pytensor/link/pytorch/linker.py b/pytensor/link/pytorch/linker.py
@@ -37,6 +37,10 @@ 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
+
         from pytensor.link.pytorch.dispatch import pytorch_typify
 
         class wrapper: