学習画像の表示の仕方がわからない

実現したいこと

https://github.com/asyml/vision-transformer-pytorch
上記サイトで独自のデータセットの画像を用意したのですが、その画像を出力表示させる方法がわかりません。
情報の不足があるかもしれませんがお願いいたします。

発生している問題・分からないこと

https://reafnex.net/ai/pytorch-use-imagefolder/
上記リンクを参考にdata_loader.pyのコードに追記しましたがエラーが出ます。

エラーメッセージ

error
1TypeError: 'iterator' object is not subscriptable

該当のソースコード

class ImageNetDataLoader(DataLoader):
    def __init__(self, data_dir, split='train', image_size=224, batch_size=16, num_workers=8):

        if split == 'train':
            transform = transforms.Compose([
                transforms.Resize((image_size, image_size)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
            ])
        else:
            transform = transforms.Compose([
                transforms.Resize((image_size, image_size)),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
            ])

        self.dataset = ImageFolder(root=os.path.join(data_dir, split), transform=transform)
        
        super(ImageNetDataLoader, self).__init__(
            dataset=self.dataset,
            batch_size=batch_size,
            shuffle=True if split == 'train' else False,
            num_workers=num_workers)
            
        imgs = iter(self.dataset)
        print("image shape ==>;",imgs[0].shape)
        pic = transforms.ToPILImage(mode='RGB')(imgs[0])
        plt.imshow(pic)
        print("Label is ",labels[0].numpy())