Bhack:

I think that you still need to trigger the epoch event in your custom training loop.


I guess for that I might need to discard the train_step() override which I don’t want to do. I will study the links you shared and get back.