My code should work, if you just change the mnist.data attribute for mnist.test_data, if you are using the test set in an older torchvision version.
What else did you change? Could you compare your code to mine?