How to create a module from graph

Hey,
I plan to export a torchscript graph represantation by using trace(...) command, and then using trace.graph(). after that, I plan to apply some optimisations for the given graph and then I want to revert the result back to ScriptModule class.
I am familiar with: torch._C._create_function_from_graph('forward', graph)
but it returns ScriptFunction class. I would like to use it as a module (ScriptModule class).

I need it as a module class because I need to be able to access its parameters by parameters() method.

thanks, Omer.

if you are mutating the graph in-place, the original ScriptModule should reflect your optimizations. Today we have no public way of creating a ScriptModule from a graph alone

Thanks for responding,
what do you mean by “in-place” in this case ?
lets assume that I have a traced module that I got by: torch.jit.trace(model, sample_inputs).
do you mean that I need to change its “graph” field or “inlined_graph” field to make the module to represent the new graph? (thats mean that all its parameters should be changed if needed)

@Michael_Suo In case we do not want to modify the graph in-place,
that is, we want an nn.Module that, given a graph, we will have the same forward function as the graph,
and will have same parameters/buffers/modules as used in the graph,

Can we automatically create such nn.Module given a graph by inferring everything needed?