Hi @Paze!
I think you might actually BE training on your GPU, but because you have your batch size set to 1, your algorithm is spending most of its time passing data to and from the GPU memory instead of actually doing the calculations. Have you tried a larger batch size?
If you’d like to see what’s taking up all of your time, you can use torch.utils.bottleneck.
https://pytorch.org/docs/1.5.1/bottleneck.html?highlight=torch%20utils%20bottleneck
Other than that, it looks like you aren’t using “with torch.no_grad():” in your evaluation routine. I’m pretty sure that you should be wrapping any inference related calls to your model/loss function inside such a statement.
Hope this helps. Good luck!
–SEH
EDIT:  Here’s another link I found extremely useful when facing a similar issue.

  
    
    
    How to prefetch data when processing with GPU? 
  
  
    I’ve spent some time working on this problem over a variety of projects. I’ve cut and pasted some past thoughts, bullets, etc from previous discussions. My background involves architecting systems that move large volumes of data from network cards to storage and then back again on request, with processing between the steps. A very similar set of concerns. 
The two main constraints that usually dominate your PyTorch training performance and ability to saturate the shiny GPUs are your total CPU IP…