Solved by ptrblck in post #2 
              
              
                I assume “overloading” means your GPU is running our of memory? 
If so, you could try to wrap this special forward pass in a with torch.no_grad() block so avoid storing the intermediate activations (which are not needed, if you don’t plan on calling backward). 
If that doesn’t help and you are still…