@slavavs I am having similar issue. Just using for loop and it seems that it is taking 4x more time to train comparing to an equivalent tensorflow model. I am struggling to find an issue. So, it would be great to see a solution if you find any.