Hi,
Could you try to do the following?:
prediction = torch.max(output.detach().cpu(), 1)[1]

As you are listing predictions, the whole computational graph is being stored (and it’s probably allocated in gpu right?)