@ptrblck , thanks for your reply,. I found that mask = (x * 0).flatten() wasted a lot of time. Then I modified my code. Now it looks like
n = torch.prod(torch.tensor(x.shape)).item()
index = np.random.randint(0, n, k)
mask = np.zeros((n,), dtype=np.float32)
mask[index] = 1
mask = mask.reshape(*x.shape)
mask = torch.from_numpy(mask)
x.mul_(mask)

The improvement in speed is obvious. I am looking for a more faster implementation…