Solved by ptrblck in post #2 
              
              
                Wrap your code in a with torch.no_grad() block, if you don’t need to call .backward() in this part of the code or detach the tensor from the computation graph, so that it can be freed via: 
predictions[i][key].detach().to('cpu')

Even though you are pushing the prediction to the CPU, the attached co…