In exercise **MedNIST Dataset**, if you run the code examples given, the code example in part **Evaluating the model** causes an error:  I am using - Chrome and - Goole Colab