Can torchscript dump backward graph?

[yxd886]

Can torchscript dump backward graph? When using TorchScript, can the backward graph be dumped and compiled using MLIR pass ? Or only forward graph is supported?

[byronyi]

Take a look at the BERT training recipe here.

It combines 2 different techniques from PyTorch core:

• torch.nn.utils.stateless.functional_call to transform model.forward(inputs) to forward(model, inputs)
• torch.fx.experimental.proxy_tensor.make_fx to transform a eager train step function into a torch.fx.GraphModule

One last thing missing would be making sure that the same train function is applied in each training step, which will require users' explicit annotation here. For more information on this, I would recommend this awesome post that summarizes the status quo of PyTorch graph capture (courtesy @penguinwu).