The first thing I’d do is to use torch.autograd.gradcheck to verify that your gradients are right.
Also, you might want to check torch.topk, which is faster than torch.sort if you only want a subset of the elements of the tensor.