JIT runtime error with nested lists


(LSM) #1

Hi, I’m getting an error when running a dummy jit script. It seems it cannot infer types on a list of lists.

import torch


class PreProcessor(torch.jit.ScriptModule):

    def __init__(self):
        super(PreProcessor, self).__init__()

    @torch.jit.script_method
    def forward(self, frames):
        # type: (List[Tensor]) -> List[Tensor]
        lidars = []
        for i in range(len(frames)):
            frame = frames[i]
            lidars.append(frame)
        return lidars


class Inference(torch.jit.ScriptModule):

    def __init__(self):
        super(Inference, self).__init__()
        self.preprocessor = PreProcessor()

    @torch.jit.script_method
    def forward(self, batched_frames):
        # type: (List[List[Tensor]]) -> List[List[Tensor]]
        data = []
        for i in range(len(batched_frames)):
            frames = batched_frames[i]
            preprocessed_data = self.preprocessor(frames)
            data.append(preprocessed_data)
        return data


if __name__ == "__main__":
    p = Inference()

    print(p)
    p.save("p.pt")

The relevant stack trace output:

RuntimeError: 
arguments for call are not valid:
  for operator aten::append(t[] self, t el) -> t[]:
  could not match type Tensor[] to t in argument 'el': type variable 't' previously matched to type Tensor is matched to type Tensor[]
  @torch.jit.script_method
  def forward(self, batched_frames):
      # type: (List[List[Tensor]]) -> List[List[Tensor]]
      data = []
      for i in range(len(batched_frames)):
          frames = batched_frames[i]
          preprocessed_data = self.preprocessor(frames)
          data.append(preprocessed_data)
                      ~~~~~~~~~~~~~~~~~ <--- HERE
      return data

Is this a known issue?


(Michael Suo) #2

Hi! Lists are statically typed in TorchScript. The default type for a list declared like foo = [] is List[Tensor]

So in data.append(preprocessed_data) you are trying to append a list of tensors to a list of tensors.

You can fix this by annotating data properly (search “Variable Type Annotation” here).