Slicing using tensor doesn't work in torchscript as expected

discort: <p>To reproduce:</p>
<p>Expected behaviour:</p>
<pre><code class="lang-auto">def func(obj_x, obj_ind):
    new_x = torch.zeros((20, 3, 7), device=obj_x.device, dtype=torch.float32)
    j = torch.zeros((1,), device=obj_x.device, dtype=torch.long)
    for i, idx in enumerate(obj_ind[1:], start=1):
        new_x[idx][j] = obj_x[i]
    return new_x

obj_x = torch.randn(10, 7)
obj_ind = torch.tensor(list(range(0, 20, 2)))
func(obj_x, obj_ind)
&gt;&gt; tensor([[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

        [[-0.2289, -1.1137,  0.2680,  0.0294,  0.3793, -0.1392, -0.7233],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 0.0367,  0.5694, -0.1834,  0.9413,  1.1738, -1.4721, -2.0817],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],
...
</code></pre>
<p>Using the torchscript:</p>
<pre><code class="lang-auto">@torch.jit.script
def func(obj_x, obj_ind):
    new_x = torch.zeros((20, 3, 7), device=obj_x.device, dtype=torch.float32)
    j = torch.zeros((1,), device=obj_x.device, dtype=torch.long)
    for i, idx in enumerate(obj_ind[1:], start=1):
        new_x[idx][j] = obj_x[i]
    return new_x

func(obj_x, obj_ind)
tensor([[[0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.]],

        [[0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.]],

        [[0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.]],

        [[0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0., 0., 0.]],
...
</code></pre>
<p>Did I miss something?</p>

ptrblck: <p>I think you are hitting multiple issues.<br>
The first one seems to be the “double indexing” which doesn’t assign the values to the original tensor.<br>
Index <code>new_x</code> and the assignment should work:</p>
<pre><code class="lang-python">new_x[idx, j] = obj_x[i]
</code></pre>
<p>However, even though <code>new_x</code> would contain values now, it seems that <code>enumerate(..., start=idx)</code> is not working correctly while scripting and the <code>start</code> value is ignored:</p>
<pre><code class="lang-python">def fun():
    for i, idx in enumerate(torch.arange(5), start=1):
        print(i, idx)
fun()

@torch.jit.script
def fun():
    for i, idx in enumerate(torch.arange(5), start=1):
        print(i, idx)
fun()
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">1 tensor(0)
2 tensor(1)
3 tensor(2)
4 tensor(3)
5 tensor(4)

0 0
[ CPULongType{} ]
1 1
[ CPULongType{} ]
2 2
[ CPULongType{} ]
3 3
[ CPULongType{} ]
4 4
[ CPULongType{} ]
</code></pre>
<p>Would you mind creating an issue for this on GitHub?</p>

discort: <p>Thanks for the response <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>At least now I know how to tackle the first issue. Also, I created the issue regarding the bug you’ve found <a href="https://github.com/pytorch/pytorch/issues/67142" class="inline-onebox" rel="noopener nofollow ugc">enumerate(..., start=idx) is not working correctly while scripting · Issue #67142 · pytorch/pytorch · GitHub</a></p>
<p>Feel free to continue our discussion on Github.</p>