How to set minimum 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 set minimum k elements in a dimension (e.g., dim=2) to a specific value (e.g., 5):
[[[0.8823, 0.9150, 5, 0.9593, 5],
  [0.6009, 5, 0.7936, 0.9408, 5],
  [0.9346, 5, 0.8694, 5, 0.7411]],

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


You can use a combination of torch.topk and torch.Tensor.scatter_.
(Bacause torch.topk return max_top_k and you want min_top_k. We can use -1*all_num for getting min_top_k)
val, ind = torch.topk(-a, k=2)
a.scatter_(index=ind, dim=-1, value=5)
print(a)


tensor([[[0.8823, 0.9150, 5.0000, 0.9593, 5.0000],
         [0.6009, 5.0000, 0.7936, 0.9408, 5.0000],
         [0.9346, 5.0000, 0.8694, 5.0000, 0.7411]],

        [[5.0000, 0.8854, 0.5739, 5.0000, 0.6274],
         [5.0000, 0.4414, 0.2969, 0.8317, 5.0000],
         [0.2695, 0.3588, 5.0000, 0.5472, 5.0000]]])


Input:
&gt;&gt;&gt; a = 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]]])


&gt;&gt;&gt; torch.topk(-a, k=2)
# values=tensor(
#     [[[-0.3829, -0.3904],
#       [-0.1332, -0.2566],
#       [-0.5677, -0.5936]],

#      [[-0.2666, -0.4294],
#       [-0.1053, -0.2696],
#       [-0.0062, -0.1994]]]),

# indices=tensor(
#     [[[2, 4],
#       [4, 1],
#       [3, 1]],

#      [[3, 0],
#       [4, 0],
#       [4, 2]]])

