Great !
By the way, unless your code specifically does not handle batch with a size different than the full one, you don’t need the if mb.size(0) != bs: block. Pytorch functions will all work with a batch size that is different and you can move this print outside of the inner loop if the goal is just to have a print !