Support a to_torchscript function on the LightningModule
🚀 Feature
Support a conversion function to PyTorch JIT similar to what's available for ONNX.
Motivation
TorchScript is a way to create serializable and optimizable models from PyTorch code. Any TorchScript program can be saved from a Python process and loaded in a process where there is no Python dependency. TorchScript is a method by which users can serve PyTorch models efficiently.
Pitch
By default, we can use TorchScript to script the LightningModule. Users can override this in their own lightning modules to use tracing, or to script specific nn.Modules inside their LightningModule. This can then be extended to other Lightning utilities like model checkpointing, so we can save TorchScript or ONNX converted models alongside the best model checkpoints to make going to serving even easier to do
    def to_torchscript(self):
        """Saves the model as a JIT module. 
            This can be overridden to support custom TorchScript module export
        Example:
            >>> class SimpleModel(LightningModule):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self.l1 = torch.nn.Linear(in_features=64, out_features=4)
            ...
            ...     def forward(self, x):
            ...         return torch.relu(self.l1(x.view(x.size(0), -1)))
            >>> with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
            ...     model = SimpleModel()
            ...     torch.jit.save(model.to_torchscript(), tmpfile.name)
            ...     os.path.isfile(tmpfile.name)
            True
        """
       return torch.jit.script(self.eval())