Hello, I am a student currently studying this code. I came across the demo_eee_CIFAR100.ipynb file in the notebooks folder, and I would like to know how to use other checkpoints to test the model's accuracy on different datasets, including how to use gs://vmoe_checkpoints/eee_s32_last2_ilsvrc2012_ft_cifar100to test on a non-CIFAR100 dataset. I often encounter issues where the input shape and the expected shape are not matching when attempting this on my own. I would greatly appreciate it if you could help me with this question!