How to remove `for loop` for this snippet?

1 
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]]])
2 
>>> 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
3

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