Ran out of memory while using torch.index_select 10
>>> print(input.size())
(1, 1, 256, 512)
>>> print(idx.size())
131072
>>> pix = torch.index_select(input, 0, idx)
RuntimeError: cuda runtime error(2): out of memory at /pytorch/torch/lib/THC/generic/THCStorage.cu:58

Any fix? Thank you.