The error is because of the following python line:
train_labels = train_labels / 255.0
Here you have divided by 255.0 to scale the values so that it reduces the computational time. This is called Normalization. You need to do normalization only for train_images but not for train_labels. So, remove that mentioned line above and it will work fine.
You can register for this AI Course by Intellipaat that teaches Artificial Intelligence in TensorFlow with Instructor-led training, good projects, and case studies.