Hello,
Your tensors might be huge and the bandwidth to your GPU not big enough thus slowing down the process.
You can reduce your batchsize and accumulate the gradient (by not using optimizer.zero_grad() at every backward() call). This will most likely reduce the bottleneck around the bandwidth.
if you reduce your batchsize by let’s say 4 you can use
if counter % 4 == 0:
    optimizer.zero_grad()
.
have fun with your loss here
.
if counter % 4 == 0:
    optimizer.step()
counter += 1