Assume proposals is of type torch.cuda.FloatTensor, and of size (some_big_num, 4). It follows that
scores, idx = torch.topk(scores, pre_nms_topN, 0, sorted=True)
proposals = proposals.index_select(0, idx)

raise the error “index_select received an invalid combination of arguments - got (int, !torch.cuda.LongTensor!), but expected (int dim, torch.LongTensor index)”. I can fix the bug with
scores, idx = torch.topk(scores, pre_nms_topN, 0, sorted=True)
proposals = proposals.index_select(0, idx.cpu())

But it incurs data exchange b/w GPU and CPU. Is there a more elegant solution? This issue is quite common when you implement some layers which are expected to run on the GPU.