From 157b0251a3501b9d301af3f80fa47c1da0470387 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 14 Jul 2024 01:06:07 +0200 Subject: [PATCH] Fix `TORCHVISION_0_18` for `allow_empty=True` (#14415) Signed-off-by: Glenn Jocher Co-authored-by: UltralyticsAssistant --- ultralytics/data/dataset.py | 4 ++-- ultralytics/utils/torch_utils.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ultralytics/data/dataset.py b/ultralytics/data/dataset.py index c10adef8d0..3ed6a21962 100644 --- a/ultralytics/data/dataset.py +++ b/ultralytics/data/dataset.py @@ -15,7 +15,7 @@ from torch.utils.data import ConcatDataset from ultralytics.utils import LOCAL_RANK, NUM_THREADS, TQDM, colorstr from ultralytics.utils.ops import resample_segments -from ultralytics.utils.torch_utils import TORCH_1_13 +from ultralytics.utils.torch_utils import TORCHVISION_0_18 from .augment import ( Compose, @@ -417,7 +417,7 @@ class ClassificationDataset: import torchvision # scope for faster 'import ultralytics' # Base class assigned as attribute rather than used as base class to allow for scoping slow torchvision import - if TORCH_1_13: # 'allow_empty' argument first introduced in torch 1.13 + if TORCHVISION_0_18: # 'allow_empty' argument first introduced in torchvision 0.18 self.base = torchvision.datasets.ImageFolder(root=root, allow_empty=True) else: self.base = torchvision.datasets.ImageFolder(root=root) diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py index ca814b602c..db17813f8f 100644 --- a/ultralytics/utils/torch_utils.py +++ b/ultralytics/utils/torch_utils.py @@ -40,6 +40,7 @@ TORCH_2_0 = check_version(torch.__version__, "2.0.0") TORCHVISION_0_10 = check_version(TORCHVISION_VERSION, "0.10.0") TORCHVISION_0_11 = check_version(TORCHVISION_VERSION, "0.11.0") TORCHVISION_0_13 = check_version(TORCHVISION_VERSION, "0.13.0") +TORCHVISION_0_18 = check_version(TORCHVISION_VERSION, "0.18.0") @contextmanager