GraphModule.to_folder produces bad syntax in module.py

torch.compile FX (Functional Transformations)
1

I’m trying to leverage some new FX features, and part of the pipeline I have in mind would involve the to_folder method. However, it produces sometimes syntactically incorrect outputs. Here’s an example:

from torch.fx import symbolic_trace
import torchvision.models as tvm

model = tvm.resnet18()
model_fx = symbolic_trace(model)
model_fx.to_folder('debug', 'Resnet18Transformed')
from debug import Resnet18Transformed

I get the following error:

    def forward(self, x : torch.Tensor) -> torch.Tensor:
    ^
IndentationError: unexpected indent

It seems like because the forward function has a dependency on the torch module (it calls torch.flatten), FX tries to insert an “import torch” near the dependency, even if that insertion breaks syntax and despite that there’s already another “import torch” in the beginning.

Any better idea on why this is happening, and what does it take to fix fundamentally? Thanks