How to process a tensor without using a for loop?

Question

    tensor = torch.zeros_like(out, dtype=torch.int).to(out.device)
    for b in range(out.shape[0]):
        for index in range(out.shape[1]):
            n = numbers[b,index]
            tensor[b,index,torch.randperm(out.shape[2])[:n]] = 1

这里tensor是一个[128,100,512]的张量，里面全为0，number是一个[128,100]的张量，里面的数字均为0~512之间。这段代码意思是需要将tensor的第三个维度随机变为1，随机变化的个数由number里的数决定。例如number[5,5]里面为96，那么tensor[5,5]中则有96个数字被随机变成0。我该怎么样将这段代码改成相率更高，因为使用for循环的速度太慢了 This tensor here is a [128, 100, 512] tensor, all filled with zeros. 'number' is a [128, 100] tensor, with numbers ranging from 0 to 512. The purpose of this code is to randomly set the third dimension of the 'tensor' to 1, and the number of random changes is determined by the numbers in the 'number' tensor. For example, if 'number[5, 5]' is 96, then 96 numbers in 'tensor[5, 5]' will be randomly changed to 1. How can I modify this code to increase the probability, as using a for loop is too slow?

我试过问ChatGPT但是它并没有给我想要的答案