Error with my custom concat class with TorchScript

jit
1

Hi I try to use torch.jit.script on my custom Concat class:

class Concat(nn.Module):
    def __init__(self, dimension=1):
        super(Concat, self).__init__()
        self.dim = dimension

    def forward(self, x):
        return torch.cat(x, dim=self.dim)

But I got the following error:

RuntimeError:
Arguments for call are not valid.
The following variants are available:
aten::cat(Tensor tensors, int dim=0) → (Tensor):
Expected a value of type ‘List[Tensor]’ for argument ‘tensors’ but instead found type ‘Tensor (inferred)’.
Inferred the value for argument ‘tensors’ to be of type ‘Tensor’ because it was not annotated with an explicit type.

aten::cat.names(Tensor tensors, str dim) → (Tensor):
Expected a value of type ‘List[Tensor]’ for argument ‘tensors’ but instead found type ‘Tensor (inferred)’.
Inferred the value for argument ‘tensors’ to be of type ‘Tensor’ because it was not annotated with an explicit type.

aten::cat.names_out(Tensor tensors, str dim, *, Tensor(a!) out) → (Tensor(a!)):
Expected a value of type ‘List[Tensor]’ for argument ‘tensors’ but instead found type ‘Tensor (inferred)’.
Inferred the value for argument ‘tensors’ to be of type ‘Tensor’ because it was not annotated with an explicit type.

aten::cat.out(Tensor tensors, int dim=0, *, Tensor(a!) out) → (Tensor(a!)):
Expected a value of type ‘List[Tensor]’ for argument ‘tensors’ but instead found type ‘Tensor (inferred)’.
Inferred the value for argument ‘tensors’ to be of type ‘Tensor’ because it was not annotated with an explicit type.

I am using PyTorch version 1.7.1+cu92, and Python 3.7.6. Is it related to this bug?

2

TorchScript is a statically-typed language, and, in most cases, we’re forced to infer untyped variables to be instances of torch.Tensor. (One of our projects this half is improving our type inference.) An easy fix is to simply annotate the relevant variables with their correct type.

def forward(self, x: List[torch.Tensor]):
    return torch.cat(x, dim=self.dim)
3 Likes
3

Thank you it is fixed by your proposal.

1 Like
4

Thank you, it solved my problem too.