In [1]: import torch

In [2]: from tsai.models.RNN import LSTM

In [3]: from tsai.models.TransformerModel import TransformerModel

In [4]: optimized_transformer = torch.jit.script(TransformerModel(c_in=2, c_out=1))
<RETRACTED>/venv/lib/python3.9/site-packages/torch/jit/_recursive.py:240: UserWarning: 'batch_first' was found in ScriptModule constants, but was not actually set in __init__. Consider removing it.
  warnings.warn("'{}' was found in ScriptModule constants, "
<RETRACTED>/venv/lib/python3.9/site-packages/torch/jit/_recursive.py:234: UserWarning: 'norm' was found in ScriptModule constants,  but it is a non-constant submodule. Consider removing it.
  warnings.warn("'{}' was found in ScriptModule constants, "

In [5]: optimized_lstm = torch.jit.script(LSTM(c_in=2, c_out=1))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 optimized_lstm = torch.jit.script(LSTM(c_in=2, c_out=1))

File <RETRACTED>/venv/lib/python3.9/site-packages/torch/jit/_script.py:1265, in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1263 if isinstance(obj, torch.nn.Module):
   1264     obj = call_prepare_scriptable_func(obj)
-> 1265     return torch.jit._recursive.create_script_module(
   1266         obj, torch.jit._recursive.infer_methods_to_compile
   1267     )
   1269 if isinstance(obj, dict):
   1270     return create_script_dict(obj)

File <RETRACTED>/venv/lib/python3.9/site-packages/torch/jit/_recursive.py:454, in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    452 if not is_tracing:
    453     AttributeTypeIsSupportedChecker().check(nn_module)
--> 454 return create_script_module_impl(nn_module, concrete_type, stubs_fn)

File <RETRACTED>/venv/lib/python3.9/site-packages/torch/jit/_recursive.py:520, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    518 # Compile methods if necessary
    519 if concrete_type not in concrete_type_store.methods_compiled:
--> 520     create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    521     # Create hooks after methods to ensure no name collisions between hooks and methods.
    522     # If done before, hooks can overshadow methods that aren't exported.
    523     create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs)

File <RETRACTED>/venv/lib/python3.9/site-packages/torch/jit/_recursive.py:371, in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    368 property_defs = [p.def_ for p in property_stubs]
    369 property_rcbs = [p.resolution_callback for p in property_stubs]
--> 371 concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)

RuntimeError:
Module 'LSTM' has no attribute 'dropout' (This function exists as an attribute on the Python module, but we failed to compile it to a TorchScript function.
The error stack is reproduced here:
Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "<RETRACTED>/venv/lib/python3.9/site-packages/fastai/imports.py", line 66
def noop (x=None, *args, **kwargs):
                          ~~~~~~~ <--- HERE
    "Do nothing"
    return x
:
  File "<RETRACTED>/venv/lib/python3.9/site-packages/tsai/models/RNN.py", line 21
        output, _ = self.rnn(x) # output from all sequence steps: [batch_size x seq_len x hidden_size * (1 + bidirectional)]
        output = output[:, -1]  # output from last sequence step : [batch_size x hidden_size * (1 + bidirectional)]
        output = self.fc(self.dropout(output))
                         ~~~~~~~~~~~~ <--- HERE
        return output