You could use this code from the ImageNet example to check, how long your training loop spends on the data loading part.
If you see a constant high number, you could use to increase / lower the number of workers etc. I assume you are already using a DataLoader
. If not, have a look at this tutorial to see, how to use one.
I also recommend to read this post for a better understanding of potential data loading bottlenecks.