@ptrblck Excuse me, now i have a tensor
x = torch.rand(1, 2, 3)

and I want to generate another random mask tensor, which has k positive elements.
The following code could solve my problem, but it is too inefficient.
n = torch.prod(torch.tensor(x.shape)).item()
index = torch.multinomial(torch.rand(n), k, replacement=False)
mask = (x * 0).flatten()
mask[index] = 1
mask = mask.reshape(*x.shape)
x.mul_(mask)

Have any suggestion?