Extract graph for editing

What I need is:

  1. to edit graph for original nn.Module model;
  2. train it
  3. to edit graph again

The only way I’ve found to achieve graph editing using pytorch2.0 tools is using custom_backend but the point that it triggers only one time so that my third case in the list above is not achievable.
What are the existing options?

This tutorial might help Google Colab, you can move from an aten graph back to an nn module easily with some of this infra

  1. is this applicable for aten graph only (not for torch graph)?
  2. if yes for first question than considering some example from the notebook
import torch._dynamo
from torch._functorch.aot_autograd import aot_module_simplified

def toy_backend(gm, sample_inputs): 
    def my_compiler(gm, sample_inputs):
        # <implement your compiler here>
        return gm.forward

    # Invoke AOTAutograd
    return aot_module_simplified(

fn = torch.compile(backend=toy_backend, dynamic=True)(model)

# triggers compilation of forward graph on the first run
out = fn(input)

# triggers compilation of backward graph on the first run

we can change graph in the first time in place1 but what about second time?