diff --git a/docs/source/examples/mae.rst b/docs/source/examples/mae.rst index eae27ed7e..73d53875c 100644 --- 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 `_ +.. note:: + + MAE requires `TIMM `_ 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 index c7cc71e71..6572313da 100644 --- 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 index 1f22ae8ec..2fffacfb9 100644 --- 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 index 832ebab17..3c29ae487 100644 --- 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 index e3b6ccfec..d4414d210 100644 --- 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):