I am just wondering why your let train_loader shuffle = False and test_loader shuffle=True? Should it be vice versa?