Как вывести на экран случайные изображения из train_data?

Я написал функцию для вывода шести рандомных изображений с лейблами из моих тренировочных данных. Лейблы она успешно выводит на экран, но почему-то не показывает изображения.

def show_transform_images(dataset):
    loader = torch.utils.data.DataLoader(dataset, batch_size = 6, shuffle = True)
    batch = next(iter(loader))
    images, labels = batch

    grid = torchvision.utils.make_grid(images, nrow=3)
    plt.figure( figsize = (11, 11) )
    plt.imshow( np.transpose(grid, (1,2,0)))
    print('labels: ', labels)

show_transform_images(train_data)