You could try printing several “intermediate” variables, like X and y, output, res, to make sure that you get different values between the train loader and test loader.
The most important thing is checking y[i] against elem, this way you will see if your model makes mistakes.
One thing I’ve noticed in your code (but it might be due to copy/pasting and not be present in your actual code) is in the first error measurement (on the test set, i.e. the 2nd code block in your message): you print res/len(dset_train). Shouldn’t it be res/len(dset_test) ?