fix import when transforms.v2 not available

lightly-ai · Nov 12, 2024 · 2e89959 · 2e89959
1 parent edd1755
commit 2e89959
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lightly/transforms/__init__.py b/lightly/transforms/__init__.py
@@ -8,7 +8,6 @@
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
 
-from lightly.transforms.add_grid_transform import AddGridTransform
 from lightly.transforms.aim_transform import AIMTransform
 from lightly.transforms.amplitude_rescale_transform import AmplitudeRescaleTranform
 from lightly.transforms.byol_transform import (
@@ -55,3 +54,7 @@
 from lightly.transforms.vicreg_transform import VICRegTransform, VICRegViewTransform
 from lightly.transforms.vicregl_transform import VICRegLTransform, VICRegLViewTransform
 from lightly.transforms.wmse_transform import WMSETransform
+from lightly.utils.dependency import torchvision_transforms_v2_available
+
+if torchvision_transforms_v2_available():
+    from lightly.transforms.add_grid_transform import AddGridTransform