Add timm version info to docs (#1517)

* Add minimum timm version info to docs * Update MAE transform import in docs
lightly-ai · Mar 22, 2024 · 158cf74 · 158cf74
1 parent 489c87c
commit 158cf74
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 4 deletions.
diff --git a/docs/source/examples/mae.rst b/docs/source/examples/mae.rst
@@ -15,6 +15,14 @@ input image and using a shallow decoder architecture.
 Reference:
     `Masked Autoencoders Are Scalable Vision Learners, 2021 <https://arxiv.org/abs/2111.06377>`_
 
+.. note::
+
+    MAE requires `TIMM <https://github.com/huggingface/pytorch-image-models>`_ to be
+    installed
+
+    .. code-block:: bash
+
+        pip install "timm>=0.9.9"
 
 .. tabs::
     .. tab:: PyTorch

diff --git a/docs/source/getting_started/install.rst b/docs/source/getting_started/install.rst
@@ -37,7 +37,7 @@ If you want to work use the Masked Autoencoder you need to additionally install
 
 .. code-block:: bash
 
-    pip install timm
+    pip install "timm>=0.9.9"
 
 Next Steps
 ------------

diff --git a/examples/pytorch/mae.py b/examples/pytorch/mae.py
@@ -8,7 +8,7 @@
 
 from lightly.models import utils
 from lightly.models.modules import MAEDecoderTIMM, MaskedVisionTransformerTIMM
-from lightly.transforms.mae_transform import MAETransform
+from lightly.transforms import MAETransform
 
 
 class MAE(nn.Module):

diff --git a/examples/pytorch_lightning/mae.py b/examples/pytorch_lightning/mae.py
@@ -9,7 +9,7 @@
 
 from lightly.models import utils
 from lightly.models.modules import MAEDecoderTIMM, MaskedVisionTransformerTIMM
-from lightly.transforms.mae_transform import MAETransform
+from lightly.transforms import MAETransform
 
 
 class MAE(pl.LightningModule):

diff --git a/examples/pytorch_lightning_distributed/mae.py b/examples/pytorch_lightning_distributed/mae.py
@@ -9,7 +9,7 @@
 
 from lightly.models import utils
 from lightly.models.modules import MAEDecoderTIMM, MaskedVisionTransformerTIMM
-from lightly.transforms.mae_transform import MAETransform
+from lightly.transforms import MAETransform
 
 
 class MAE(pl.LightningModule):