Export torch script when training, Allow to import it for inference

Describe the solution you'd like
As user trains models, we are changing a lot of hyperparameters which includes number or layer in the network architecture and even change the model architecture completely. This requires user to track the model weight and architecture pairs in the inference.

Describe alternatives you've considered
We can export ts next to checkpoint models which holds the model weights

Additional context

• Exporting ts can affect performance so it should be a done at the end of the training.
• Need to accommodate for user stopping the training

No files, tests, or entry points are named in the issue. First clarify where training checkpoints are created and how inference loads models, then determine how export timing should handle performance and interrupted training. Done should include a defined TorchScript export and a documented way to import it for inference.