Hello @Linux-cpp-lisp,
This is a good question. We didn’t foresee the desire to do this, so the current APIs aren’t factored in a clean enough way to easily support something like this. (You can override fx.Graph.python_code, but it would involve a lot of copy-paste, which might be OK for some situations). This seems like a good idea, though, and I’ve filed a feature request: [FX] Allow customizing code generation for certain Nodes · Issue #53732 · pytorch/pytorch · GitHub 2
For issue #53517 in particular, the TorchScript fix is being rolled up into the 1.8.1 patch release, with release date to be announced