How to modify a torch ExportedProgram with custom operator by torch.export?

1

I train a model with a custom gpu layer implementation, let’s call it as CustomLookupModule. Now I want to export the model for serving in C++. But I’ve met a few issues.

  1. The export method needs to run IR decomposition for unsupported operators, how can I export a module as a whole, like an standalone operator called CustomLookupModule?
  2. during serving/inference, the module has different implementation from training. So I need to modify the ExportedProgram to replace the operator into custom implementation. How can I do it?
  3. How to I make AOT compilation support custom operator for ExportedProgram?

Thanks.