torch.topk (torch.topk — PyTorch 1.10.0 documentation) should get you the k smallest elements along with their indices if you use the argument largest=False. Should be much faster than sorting the whole Tensor. You can then use those indices to do what you want.