Bem observado, Vinícius!
O certo seria ter um loop para percorrer os batches na base MNIST também, pois na aula é utilizado somente 1 batch (128 registros) em cada época, mas existem mais dados disponíveis para treinamento. O script apresentado é mais para demonstrar o funcionamento do algoritmo e da pipeline, mas certamente pode ser melhorado.