Fix comment: https://github.com/rwightman/pytorch-image-models/pull/1564#issuecomment-1326743424

2 years ago · 653bdc7105
parent bdc9fad638
commit 653bdc7105
1 changed files with 5 additions and 0 deletions
--- a/timm/data/dataset_factory.py
+++ b/timm/data/dataset_factory.py
@ -131,7 +131,12 @@ def create_dataset(
            elif split in _EVAL_SYNONYM:
                split = 'val'
            ds = Places365(split=split, **torch_kwargs)
+        elif name == 'qmnist':
+            assert has_qmnist, 'Please update to a newer PyTorch and torchvision for QMNIST dataset.'
+            use_train = split in _TRAIN_SYNONYM
+            ds = QMNIST(train=use_train, **torch_kwargs)
        elif name == 'imagenet':
+            assert has_imagenet, 'Please update to a newer PyTorch and torchvision for ImageNet dataset.'
            if split in _EVAL_SYNONYM:
                split = 'val'
            ds = ImageNet(split=split, **torch_kwargs)