Hi,
The most likely reason is that the for loop that you use in your loss function is slow and creates many small operations in the backward as well.
Maybe you want to change it to work with the whole batch at once?