IndexError: tensors used as indices must be long, byte or bool tensors

I am getting this error only during the testing phase, but I do not face any problem in the training and validation phase.

IndexError: tensors used as indices must be long, byte or bool tensors

I get this error for the last line in the given code snippet.

The code snippet looks like the one below,

NumClass = 10
mask = torch.zeros(batch_size, self.mem_dim, 4, 4)
ones = torch.ones(1, 4, 4)
NumRows = self.mem_dim
Elements = NumRows//NumClass
for i in range(batch_size):
    lab = torch.arange(Elements * label[i], Elements*(label[i]+1), 1)
    mask[i,lab] = ones

The “lab” is a tensor value and prints out the range in such a way,

tensor([6, 7, 8])
tensor([ 9, 10, 11])
tensor([21, 22, 23])

(Note*: the length of this lab tensor can be of length ‘n’ based on the value of ElementsPerClass)

Answer

Yes it works when I provide a dtype=long to the label tensor and leaving rest of the tensors with default dtypes. thanks!