I have the following code in an extension of torch.autograd.Function:

self.avarwic_infor = {
c: {
‘ind_c’: ind_range[y==c],
‘mask_c’: y==c,
‘prob_c’: torch.FloatTensor([torch.sum(y==c)]) / torch.FloatTensor([self.N])
} for c in torch.range(0, self.C-1).type(‘torch.ByteTensor’)
}

where y has type LongTensor, self.N and self.C have type long.
it works fine if I use CPU, but got the following error if GPU is used:

TypeError: indexing a tensor with an object of type ByteTensor. The only supported types are integers, slices, numpy scalars and torch.ByteTensor.

is this error because of I am using dictionary and GPU simultaneously or I did something else wrong?
Thank you very much!