Thanks for the details. I indeed think its surprising it takes so long on your GPU. It might be because we are still in the regime where the overhead is not really worth it.
As for my initial message : never mind, I was comparing to simulations which I was running with a lighter version of the code that slightly deteriorates the quality of the training. My normal code takes a similar amount as yours (though I currently run it on my CPU).
I hope that google colab solved your issue !
As for my initial message : never mind, I was comparing to simulations which I was running with a lighter version of the code that slightly deteriorates the quality of the training. My normal code takes a similar amount as yours (though I currently run it on my CPU).
I hope that google colab solved your issue !
