Merge pull request #1414 from dedeswim/bits_and_tpu

Fix issue with `torchvision`'s `ImageNet`
pull/804/head
Ross Wightman 2 years ago committed by GitHub
commit b4ea69c9ce
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -122,12 +122,14 @@ def create_dataset(
elif name == 'imagenet': elif name == 'imagenet':
if split in _EVAL_SYNONYM: if split in _EVAL_SYNONYM:
split = 'val' split = 'val'
torch_kwargs.pop("download")
ds = ImageNet(split=split, **torch_kwargs) ds = ImageNet(split=split, **torch_kwargs)
elif name == 'image_folder' or name == 'folder': elif name == 'image_folder' or name == 'folder':
# in case torchvision ImageFolder is preferred over timm ImageDataset for some reason # in case torchvision ImageFolder is preferred over timm ImageDataset for some reason
if search_split and os.path.isdir(root): if search_split and os.path.isdir(root):
# look for split specific sub-folder in root # look for split specific sub-folder in root
root = _search_split(root, split) root = _search_split(root, split)
torch_kwargs.pop("download")
ds = ImageFolder(root, **kwargs) ds = ImageFolder(root, **kwargs)
else: else:
assert False, f"Unknown torchvision dataset {name}" assert False, f"Unknown torchvision dataset {name}"

Loading…
Cancel
Save