You have train as     for batch_idx, (data, target) in enumerate(train_loader):
and test as for data, target in test_loader:. Fix that and it should work?
Also shuffling test is not a good idea, you want to keep it same for comparision across runs