diff --git a/pytensor/link/pytorch/linker.py b/pytensor/link/pytorch/linker.py index d47aa43dda..0d2e190ddc 100644 --- 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: