Conversion of torch GraphModule back to nn.Module

I’ve done some changes with my nn.Module using torch.fx and after I did recompile() call I eventually got GraphModule object.
So the question is how to convert it back to usual nn.Module or it’s better to leave it in this form cause otherwise my changes won’t be applied or smth?

You can use the GraphModule object obtained after applying torch.fx and recompile() like a regular nn.Module, as it is a subclass of nn.Module. There is no need to convert it back to the original nn.Module. The changes you made using torch.fx are preserved within the GraphModule, and you can use it for training, evaluation, and other operations just like you would with a regular nn.Module.


thanks a lot, have a good day or something)

1 Like