At the end of the "Key Attributes" section, just before figure 2.2, it is written that we display the fourth element, but it's the fifth:

...
Let’s display the fourth digit in this rank-3 tensor, using the library Matplotlib (part of the standard scientific Python suite); see figure 2.2.
...
digit = train_images[4]
...