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

実現したいこと

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())

試したこと・調べたこと

上記の詳細・結果

解決策につなが幼なものは見つけられませんでした。

補足

特になし