0

I have a function that makes a list and appends to it. I want to convert my function into a Pytorch jit function to both speed up the computation as well as fill the lists with parameters that I will eventually optimize over. I am not sure if lists are compatible with Pytorch jit functions and I am getting errors when I try to do simple examples.

For example I tried doing this

import torch

@torch.jit.script
def my_function(x):
    my_list = []
    for i in range(int(x)):
        my_list.append(i)
    return my_list

a = my_function(10)
print(a)

but I got this error

aten::append.t(t[](a!) self, t(c -> *) el) -> t[](a!):
Could not match type int to t in argument 'el': Type variable 't' previously matched to type Tensor is matched to type int.
:
  File "myscript.py", line 18
    my_list = []
    for i in range(int(x)):
        my_list.append(i)
        ~~~~~~~~~~~~~~ <--- HERE
    return my_list

What is wrong here? Am I not allowed to use lists in PyTorch? If not, what other append-able object can I substitute that is compatible with PyTorch?

Shep Bryan
  • 567
  • 5
  • 13

1 Answers1

2

Your problem is that Pytorch's JIT does not support list which are mutable.
You should convert your list to a Pytorch's tensor and then use torch.cat to concatenate so that the JIT compiler knows what is going on.

Toyo
  • 667
  • 1
  • 5
  • 22