How to randomly set k elements in a dimension in pytorch tensor to a specific value?
For example, if I have tensor (shape [2, 3, 5])
[[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904],
  [0.6009, 0.2566, 0.7936, 0.9408, 0.1332],
  [0.9346, 0.5936, 0.8694, 0.5677, 0.7411]],

 [[0.4294, 0.8854, 0.5739, 0.2666, 0.6274],
  [0.2696, 0.4414, 0.2969, 0.8317, 0.1053],
  [0.2695, 0.3588, 0.1994, 0.5472, 0.0062]]]

and k = 2, I want to randomly set k elements in a dimension (e.g., dim=2) to a specific value (e.g., 5):
[[[0.8823, 0.9150, 0.3829, 5, 5],
  [0.6009, 0.2566, 5, 0.9408, 5],
  [5, 5, 0.8694, 0.5677, 0.7411]],

 [[5, 0.8854, 0.5739, 5, 0.6274],
  [5, 0.4414, 0.2969, 5, 0.1053],
  [0.2695, 0.3588, 5, 0.5472, 5]]]


Create a random binary mask with k True elements in dimension dim using argsortand set those to 5.
import torch

p = torch.tensor(
    [[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904],
     [0.6009, 0.2566, 0.7936, 0.9408, 0.1332],
     [0.9346, 0.5936, 0.8694, 0.5677, 0.7411]],

    [[0.4294, 0.8854, 0.5739, 0.2666, 0.6274],
     [0.2696, 0.4414, 0.2969, 0.8317, 0.1053],
     [0.2695, 0.3588, 0.1994, 0.5472, 0.0062]]], requires_grad=True)

k = 2
mask = torch.rand(p.shape).argsort(2) &lt; k
torch.where(mask, 5, p)

Output
tensor([[[5.0000, 0.9150, 5.0000, 0.9593, 0.3904],
         [5.0000, 0.2566, 0.7936, 5.0000, 0.1332],
         [0.9346, 5.0000, 0.8694, 0.5677, 5.0000]],

        [[5.0000, 0.8854, 0.5739, 0.2666, 5.0000],
         [5.0000, 0.4414, 0.2969, 5.0000, 0.1053],
         [5.0000, 5.0000, 0.1994, 0.5472, 0.0062]]], grad_fn=&lt;WhereBackward0&gt;)

