GitLab

* when k=1, it reduces to torch.max, and not surprising that torch.max is
faster than torch.topk.
* however when k=2, it is even slower than torch.topk