This correspond a question I asked on Pytorch forum. When we want to use indexing to extract data which is already on GPU, should indexing arrays better be on GPU as well? The answer is yes. Here is the evidence:
I also created some other examples to show that if you are generating indexing arrays on the fly, they should be best created using torch.xxx(..., device=torch.device("cuda") rather than torch.xxx(...).
