Jit for tolist()

jit
2

This is a known bug https://github.com/pytorch/pytorch/issues/26752

You can workaround by implementing it yourself, but it can be cumbersome due to the reasons explained in the bug report (plus the fact that TorchScript does not support generic functions or recursion):

def my_1d_tolist(x):
    result: List[float] = []
    for i in x:
        result.append(i.item())
    return result

@torch.jit.script
def my_2d_tolist(x):
    result: List[List[float]] = []
    for i in x:
        result.append(my_1d_tolist(i))
    return result

x = torch.ones(2, 2)

# These calls should be the same
print(x.tolist())
print(my_2d_tolist(x))
1 Like