At this stage it might be useful to drill down on the distribution of how time is spent in the training loop by adding a bunch of time.time() statements and narrowing down what the relative cost of each operation is. Note that to do this you would want to add torch.cuda.synchronize before and after starting timing for parts containing GPU operations to ensure that the timing information is accurate.
Once you’ve optimized the bottlenecks, you would then want to remove these synchronize calls to reduce the overhead.