Torchscript error

AI228: <p>Trying to modify the below code to torch.jit.script - but facing errors…any pointers?</p>
<pre><code class="lang-auto">class WrappedLSTM(nn.Module):
    def __init__(self, d_input, d_model, n_layer, batch_first, bidirectional, bias, dropout):
        super(WrappedLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size=d_input, hidden_size=d_model, num_layers=n_layer, batch_first=batch_first,
                        bidirectional=bidirectional, bias=bias, dropout=dropout)

    #def forward(self, input, hx=None):
    #    return self.lstm(input, hx)
    @script_method
    def forward(self, input: torch.Tensor, hx: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
        # Ensure to use the specific methods of self.lstm (e.g., lstm.forward) if needed
        return self.lstm(input, hx)
</code></pre>
<pre><code class="lang-auto">class Encoder(ScriptModule):
    def __init__(self, d_input, d_model, n_layer, unidirect=False,
                 dropout=0.2, dropconnect=0., time_ds=1, use_cnn=False, freq_kn=3, freq_std=2, pack=False):
        super().__init__()

        self.time_ds = time_ds

        if use_cnn:
            cnn = [nn.Conv2d(1, 32, kernel_size=(3, freq_kn), stride=(2, freq_std)), nn.ReLU(),
                   nn.Conv2d(32, 32, kernel_size=(3, freq_kn), stride=(2, freq_std)), nn.ReLU()]
            self.cnn = nn.Sequential(*cnn)
            d_input = ((((d_input - freq_kn) // freq_std + 1) - freq_kn) // freq_std + 1) * 32
        else:
            self.cnn = None

        # For TorchScript compatibility, use a ModuleList to store layers
        self.rnn = nn.ModuleList([WrappedLSTM(d_input=d_input, d_model=d_model, n_layer=n_layer, batch_first=True,
                                             bidirectional=(not unidirect), bias=True, dropout=dropout)])

        self.unidirect = unidirect
        self.pack = pack

    @script_method
    def rnn_fwd(self, seq, hid):
        # Note: WrappedLSTM should also be TorchScript-compatible
        seq, hid = self.rnn[0](seq, hid)
        return seq, hid

    @script_method
    def forward(self, seq, hid=None):
        if self.time_ds &gt; 1:
            ds = self.time_ds
            l = ((seq.size(1) - 3) // ds) * ds
            seq = seq[:, :l, :]
            seq = seq.view(seq.size(0), -1, seq.size(2) * ds)

        if self.cnn is not None:
            seq = self.cnn(seq.unsqueeze(1))
            seq = seq.permute(0, 2, 1, 3).contiguous()
            seq = seq.view(seq.size(0), seq.size(1), -1)

        seq, hid = self.rnn_fwd(seq, hid)

        if not self.unidirect:
            hidden_size = seq.size(2) // 2
            seq = seq[:, :, :hidden_size] + seq[:, :, hidden_size:]

        return seq, hid
</code></pre>
<h1><a name="error-message-1" class="anchor" href="#error-message-1"></a>Error Message:</h1>
<p>forward(<strong>torch</strong>.eval_ipex_cnn_true.WrappedLSTM self, Tensor input, (Tensor, Tensor)? hx=None) → ((Tensor, (Tensor, Tensor))):<br>
Expected a value of type ‘Optional[Tuple[Tensor, Tensor]]’ for argument ‘hx’ but instead found type ‘Tensor (inferred)’.<br>
Inferred the value for argument ‘hx’ to be of type ‘Tensor’ because it was not annotated with an explicit type.<br>
:<br>
File “/home/psakhamo/zoom_quantization/eval_ipex_cnn_true.py”, line 115<br>
def rnn_fwd(self, seq, hid):<br>
# Note: WrappedLSTM should also be TorchScript-compatible<br>
seq, hid = self.rnn[0](seq, hid)<br>
~~~~~~~~~~~ &lt;— HERE<br>
return seq, hid<br>
‘Encoder.rnn_fwd’ is being compiled since it was called from ‘Encoder.forward’<br>
File “/home/psakhamo/zoom_quantization/eval_ipex_cnn_true.py”, line 131<br>
seq = seq.view(seq.size(0), seq.size(1), -1)</p>
<pre><code>    seq, hid = self.rnn_fwd(seq, hid)
                                 ~~~ &lt;--- HERE

    if not self.unidirect:
</code></pre>

AI228: <p>Able to get rid of above error - with below changes</p>
<pre><code class="lang-auto">    @script_method
    def rnn_fwd(self, seq,  hid:Optional[Tuple[torch.Tensor, torch.Tensor]]):
        # Note: WrappedLSTM should also be TorchScript-compatible
        seq, hid = self.rnn[0](seq, hid)
        return seq, hid

    @script_method
    def forward(self, seq,  hid:Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
        if self.time_ds &gt; 1:
            ds = self.time_ds
            l = ((seq.size(1) - 3) // ds) * ds
            seq = seq[:, :l, :]
            seq = seq.view(seq.size(0), -1, seq.size(2) * ds)
</code></pre>
<p>Seeing a new one</p>
<p>RuntimeError: Tried to serialize object <strong>torch</strong>.eval_ipex_cnn_true.Encoder which does not have a <strong>getstate</strong> method defined!</p>
<p>Added “<strong>getstate</strong>” implementation to Encoder class as below</p>
<pre><code class="lang-auto">    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove the 'pack' key from the state dictionary
        state.pop('pack', None)
        return state
</code></pre>
<p>Now - when trying to export/serialize the pt mode</p>
<pre><code class="lang-auto">torch.save(model.encoder, torch_enc_model_path)
</code></pre>
<ul>
<li>still getting below error</li>
</ul>
<p>RuntimeError: Tried to serialize object <strong>torch</strong>.eval_ipex_cnn_true.Encoder which does not have a <strong>getstate</strong> method defined!</p>