[TorchScript] PackedSequence issue

I want to export LSTM based pytorch model to TorchScript script but I failed because of such error:

RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  forward__0(__torch__.torch.nn.modules.rnn.LSTM self, Tensor input, (Tensor, Tensor)? hx=None) -> ((Tensor, (Tensor, Tensor))):
  Expected a value of type 'Tensor' for argument 'input' but instead found type '__torch__.torch.nn.utils.rnn.PackedSequence'.
  
  forward__1(__torch__.torch.nn.modules.rnn.LSTM self, __torch__.torch.nn.utils.rnn.PackedSequence input, (Tensor, Tensor)? hx=None) -> ((__torch__.torch.nn.utils.rnn.PackedSequence, (Tensor, Tensor))):
  Expected a value of type 'Optional[Tuple[Tensor, Tensor]]' for argument 'hx' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'hx' to be of type 'Tensor' because it was not annotated with an explicit type.

The original call is:
  File "", line 24
                                 enforce_sorted=False)

        output, feat = self.lstm(x, feat)
                       ~~~~~~~~~ <--- HERE
        output = pad_packed_sequence(output,
                                     batch_first=True,

Script to reproduce:

import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence


class LSTMBlock(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size

        self.lstm = nn.LSTM(input_size,
                            hidden_size,
                            batch_first=True,
                            bidirectional=False)

    def forward(self, x, input_lengths, feat=None):
        total_length = x.shape[1]
        x = pack_padded_sequence(x,
                                 input_lengths,
                                 batch_first=True,
                                 enforce_sorted=False)

        output, feat = self.lstm(x, feat)
        output = pad_packed_sequence(output,
                                     batch_first=True,
                                     total_length=total_length)[0]

        return output, input_lengths, feat

model = LSTMBlock(100, 100)
script = torch.jit.script(model)

Is there a way how to handle models trained with packed_sequence?

Did you solve this issue? I met a very similar issue when I was trying to convert LSTM pytorch model to TorchScript.

I had to change model architecture.

Hi @Bartlomiej_Roszak I’m also facing same issue.
Can you please share what changes did you make exactly?