indexing_bug.py

from __future__ import print_function
import torch

A = torch.rand(4)
idx = torch.LongTensor([0, 3])

An = A.numpy()
idxn = idx.numpy()

# Numpy indexing.
print('numpy')
print(An[idxn].shape)
print(An[idxn, ...].shape)
if A.dim() > 1:
	print(An[:, idxn].shape)
	print(An[:, idxn, ...].shape)

# Torch indexing.
print('torch')
print(A[idx].size())
print(A[idx, ...].size())
if A.dim() > 1:
	print(A[:, idx].size())
	print(A[:, idx, ...].size())

# Output.
"""
For 2 dimensional A.
# numpy
# (2, 4)
# (2, 4)
# (4, 2)
# (4, 2)
# torch
# (2L, 4L)
# (2L, 4L)
# (4L, 2L)
# Traceback (most recent call last):
#   File "bug.py", line 22, in <module>
#     print(A[:, idx, ...].size())
# TypeError: indexing a tensor with an object of type torch.LongTensor. The only supported types are integers, slices, numpy scalars and torch.LongTensor or torch.ByteTensor as the only argument.
"""