I am training a Densely Connected U-Net model on CT scan data of dimension 512x512 for segmentation task.
My network training was very slow, so I tried to profile the different steps in my code and found the optimizer.step() line to be the bottleneck. It is extremely slow and takes nearly 0.35 secs every iteration. The time taken by the other steps is as follows: 
Screenshot_20181225_170853.png303×532 23.9 KB
 .
My optimizer declaration is: optimizer = optim.Adam(model.parameters(), lr=0.001)
I cannot understand what is the reason. Can someone suggest some possible causes?
Thanks