Expand to_torchscript to support also TorchScript's trace method
🚀 Feature

Allow for the user to easily choose between TorchScript's script or trace method to create a module.
Motivation

While TorchScript's script method will work for simple models, it will not always work out of the box when models rely on Python variables to be set. This requires the user to manually annotate the model to not run into issues with script().
TorchScript's trace method on the other hand creates a traced module that is determined by running a Tensor through the network and tracks what happens during this process.
This always works, but loses design choices if present in the model.
Both script and trace have their use cases, and with a minimal extension of this function, both methods can be used.
Pitch


Add a method argument that can be set to either script or trace (default to script, which results in the current behaviour).
Add a example_inputs argument that defaults to None and can be set to any Tensor. If None is provided, this function will automatically try to use self.example_input_array. The example input is automatically send to the correct device.
Note: example_inputs's name cannot be changed as this is the name of the argument trace() expects. If named otherwise, there can be a conflict with kwargs.

This change should not break any older scripts, as it by defaults uses script.
Alternatives
Make no change and require the user to overwrite this function to use trace.
Additional context

Please assign me for this request.