Hi,
I’m looking for a way to write the graph (preferably torch.jit.script module) forward and backward into a text format file.
A simple way to write the forward:
module = torch.jit.script(model)    #  JIT for model
print(module.code)
But, is there something similar to write the backward?
The reason I’m looking for that is: I want to evaluate performance of the forward and backward for a theoretical HW.
I tried to use pytorch profile, but it does not give all the information. Does not give the graph connectivity and tensor sizes
ONNX export support only forward.
Any other idea?