torch.nn.Embedding/torch.nn.LSTMCell and torch.jit

1

Looks like torch.jit is interpreting nn.Embedding and nn.LSTMCell as python functions. The code and error are below.

import torch
class Model(torch.jit.ScriptModule):
    def __init__(self):
        super(Model, self).__init__()

        self.embedding = torch.nn.Embedding(128, 512)

    @torch.jit.script_method
    def forward(self, x):
        return self.embedding(x)

m = Model()
m.save('temp.pt')

Traceback (most recent call last):
  File "test.py", line 16, in <module>
    m.save('temp.pt')
RuntimeError: Couldn't export Python operator <python_value>

Defined at:
@torch.jit.script_method
def forward(self, x):
    return self.embedding(x)
           ~~~~~~~~~~~~~~ <--- HERE
2

Torch Script supports a subset of the builtin tensor and neural network functions that PyTorch provides. Most methods on Tensor as well as functions in the torch namespace, all functions in torch.nn.functional and all modules from torch.nn are supported in Torch Script, excluding those in the table below. For unsupported modules, we suggest using torch.jit.trace() .

Unsupported torch.nn Modules

torch.nn.modules.adaptive.AdaptiveLogSoftmaxWithLoss torch.nn.modules.normalization.CrossMapLRN2d torch.nn.modules.fold.Fold torch.nn.modules.fold.Unfold torch.nn.modules.rnn.GRU torch.nn.modules.rnn.LSTM torch.nn.modules.rnn.RNN torch.nn.modules.rnn.GRUCell torch.nn.modules.rnn.LSTMCell torch.nn.modules.rnn.RNNCell

1 Like
3

@jotline hello i face the same problem that said that the LSTMCell is not supported.

That means I can’t convert to torchscript model if I have this LSTMCell module?

I’m trying to create a variable length OCR model.