Minhas imagens n vieram em escala de cinza

Olá João! Tudo bem com você?

Parece que houve uma alteração nas cores de apresentação do plt.imshow, que atualmente utiliza 'viridis' no parâmetro cmap para apresentar as cores.

A documentação nos fornece uma dica para o caso de querermos utilizar a escala de cinza:

Para exibir uma escala de cinza, configure o mapeamento de cores usando os parâmetros cmap='gray', vmin=0, vmax=255.

Portanto, basta alterarmos o código da seguinte forma:

import matplotlib.pyplot as plt
plt.imshow(imagens_treino[0], cmap='gray', vmin=0, vmax=255)

Dessa forma, seu gráfico estará em escala de cinza.

Bons estudos.