Torch.masked_select result in out of memory when using gpu

Here are my source codes:
def forward(self, prob, target, reward):
prob: (N, C), torch Variable
target : (N, ), torch Variable
reward : (N, ), torch Variable
N = target.size(0)
C = prob.size(1)
one_hot = torch.zeros((N, C))
if prob.is_cuda:
one_hot = one_hot.cuda()
one_hot.scatter_(1,,1)), 1)
one_hot = one_hot.type(torch.ByteTensor)
one_hot = Variable(one_hot)
if prob.is_cuda:
one_hot = one_hot.cuda()
loss = torch.masked_select(prob, one_hot)
loss = loss * reward
loss = -torch.sum(loss)
return loss
When N=2000 and C=240755, once the line “loss = torch.masked_select(prob, one_hot)” runs, the memory of GPU seems doubled. Is this phenomena normal and how to prevent it?

Hi, I also encountered the same problem, could you please let me know how to solve it?

Hi, I have replaced the line “loss = torch.masked_select(prob, one_hot)” by “loss = torch.sum(prob.mul(one_hot), dim=1)”. It seems work.