Error with my custom concat class with TorchScript

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?

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)
1 Like

Thank you it is fixed by your proposal.