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?