As the title says, is there a way to return the derivative of a nn module with respect to one of its input as a separate nn module where the operations in the nn module are known?

For instance, could the graph generated by grad() be converted to a nn module that doesn’t rely on grad inside the module? I am interested in training a nn model using grad(create_graph=true), then have the structure of its derivative with respect to an input to implement the same math operations with obtained parameters on another program.