How to prevent fx to trace into my custom defined module

torch.compile FX (Functional Transformations)
1

I have a CustomConv2d module which subclasses nn.Conv2d. When I run fx.symbolic_trace, it does not stop at CustomConv2d, but will rather trace into it and generate a graph with conv_weights, F.conv2d, etc.

How to inform fx to treat CustomConv2d as an end node?

1 Like
2

One option may be to subclass Tracer and implement a new is_leaf_module method. Wondering if there is any other approaches

3

https://pytorch.org/docs/stable/fx.html#torch.fx.wrap