Hello,
I would expect your results to improve if you shuffle your training data accross classes.
To do this elegantly it is likely a good idea to define a torch.utils.data.Dataset for your dataset and then use a DataLoader (with shuffle=True).
One epoch commonly is considered be a loop over the entire dataset (all classes) so that loop should be the outermost.
Best regards
Thomas