I would try to reduce the code until the slowdown isn’t observed anymore in order to isolate the issue, as I cannot see any obvious errors.
I.e. remove all metric calculation, loss prints etc. and check the actual training alone.



 sbhatti:

and tried adding torch.cuda.empty_cache() which unfortunately did not work.


torch.cuda.empty_cache() will not avoid out of memory issues, but will release the unused device memory stored in the cache, which will slow down your code the memory needs to be reallocated again.