Skip to content

Commit c5f26fd

Browse files
author
Ian Schweer
committed
Add compiler flags to help hint more to torch
1 parent 70634f1 commit c5f26fd

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

pytensor/link/pytorch/linker.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ def conversion_func_register(*args, **kwargs):
3737
def jit_compile(self, fn):
3838
import torch
3939

40+
# two flags that tend to help our graphs
41+
torch._dynamo.config.capture_func_transforms = True
42+
torch._dynamo.config.capture_scalar_outputs = True
43+
4044
from pytensor.link.pytorch.dispatch import pytorch_typify
4145

4246
class wrapper:

0 commit comments

Comments
 (0)