using torch script

Author: Dpbm

train.py

@@ -381,17 +381,18 @@ def one_epoch(
     return last_loss
 
 
-def get_model(device: Device, state: Optional[StateDict] = None):
+def get_model(device: Device, state: Optional[StateDict] = None) -> torch.jit.ScriptModule:
     """
-    Prepare model.
+    Prepare script model (JIT).
     It creates a model, load into a given device and load weights if a state
     was provided.
     """
     model = Model().to(device)
     if state:
         model.load_state_dict(state)
 
-    return model
+    script_model = torch.jit.script(model)
+    return script_model
 
 
 def train(