Creating a TorchScript model with loss and optimizer included


Is there a recommended way to create a TorchScript model with torch.jit.trace/script that can include the loss and optimizer methods? I see the usual way is to perform a torch.jit.trace/script over the model, and the loss and optimizer methods are kept outside. My intention is to create a graph that can encapsulate as much as possible, so that each op in the graph can go via the graph optimizations and can be executed by the same runtime executor.

If there isn’t any possible way to do this, then the follow-up question is - what is the impact on the training time due to the fact that these methods become separate ops outside the purview of the TorchScript model.