You could use scatter_ as shown here:
# setup
index = torch.tensor([[0, 0, 2, 4], [0, 0, 0, 3]])
mask = torch.zeros(2, 4, 6)

# loop
for i in range(mask.size(0)):
    for j in range(mask.size(1)):
        mask[i,j,index[i,j]] = 1

# scatter
mask2 = torch.zeros(2, 4, 6)
mask2.scatter_(2, index.unsqueeze(2), 1)

print((mask == mask2).all())
> tensor(True)