Как отсортировать данные для обучения нейросети

Я начал изучать нейросети на python3 и пытаюсь написать простую нейронку для распознавание цифр. Данные для обучения скачал с THE MNIST DATABASE.

Доставать случайную цифру я могу так:

from mnist import MNIST
import random

mndata = MNIST('samples')
images, labels = mndata.load_training()

index = random.randrange(0, len(images))
print(mndata.display(images[index]))

Но так я получаю случайную цифру. Можно ли как-то получать конкретную цифру?