Hi,
Due to some reason, I need to do the operation below.

labels = torch.tensor([[1, 0, 0, 1], [1,0, 1, 1]])
counts = torch.tensor([[2, 2], [1, 3]])
F = torch.zeros(2, 2, 4)
for i in range(2):
    for j in range(2):
        F[i, j, labels[i]==j] = counts[i][j]
print(F)

Out:
tensor([[[0., 2., 2., 0.],
         [2., 0., 0., 2.]],

        [[0., 1., 0., 0.],
         [3., 0., 3., 3.]]])

The double for loop lead to the slow operation. Is there any torch function can make it faster?
Thanks!