Script `zip` function

jit
1

Hi, I wanted to know if there is a way to script a function that uses zip, specifically I have a module that interleaves two lists and it is raising an error. To reproduce:

from typing import List

import torch

def interleave(a: List[int], b: List[int]):
    return sum(zip(a, b), ())

torch.jit.script(interleave)

This raises a RuntimeError with the message iterabletree cannot be used as a value. Is there another way to do this?

2

Do you have an example that is closer to what you want to achieve? The above example can be done by converting the two List[int]s into Tensors, then doing a cat, reshape, then sum.