That’s a generic PyTorch function in that case. We’ve implemented our Trainer class to help users easily train their model and leverage GPUs/TPUs and we’re happy to help if you encounter some bugs with it but if you want to build your own loop, you should try the PyTorch forum.
3 Likes