# How to remove `for loop` for this snippet?

``````bz=2,
h = 5
w = 5
a = torch.zeros(bz, h, w)

idx_1 = torch.tensor([[0,2,1],[2,3,0]])
idx_2 = torch.tensor([[1,2,3],[3,3,4]])

# Is it possible to remove `for loop` here ?
for i in range(bz):
a[i][idx_1[i], idx_2[i]] = 1

``````

Expected results:

``````tensor([[[0,1,0,0,0],
[0,0,0,1,0],
[0,0,1,0,0],
[0,0,0,0,0],
[0,0,0,0,0]],

[[0,0,0,0,1],
[0,0,0,0,0],
[0,0,0,1,0],
[0,0,0,1,0],
[0,0,0,0,0]]])
``````
``````>>> import torch
>>> bz=2; h=5; w = 5;
>>> a = torch.zeros(bz, h, w)
>>> idx_1 = torch.tensor([[0,2,1],[2,3,0]])
>>> idx_2 = torch.tensor([[1,2,3],[3,3,4]])
>>> idx_b = torch.arange(bz).long().unsqueeze(1).expand(bz,3)
>>> idx = idx_b * h*w + idx_1 * w + idx_2
>>> b = a.view(bz*h*w)
>>> b[idx] = 1
>>> a
tensor([[[0., 1., 0., 0., 0.],
[0., 0., 0., 1., 0.],
[0., 0., 1., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]],

[[0., 0., 0., 0., 1.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 1., 0.],
[0., 0., 0., 1., 0.],
[0., 0., 0., 0., 0.]]])
``````

This code can be made simple by using index_put.

``````a[(idx_b,idx_1,idx_2)] = 1
``````

Thank you for your help！ It is a nice solution!