From c97112e7f5ce58e8eeb72d5bb8cc9fb0e84b556a Mon Sep 17 00:00:00 2001
From: Ersi <ersi@lightly.ai>
Date: Fri, 5 Jan 2024 09:45:40 +0000
Subject: [PATCH] fixed import

---
 examples/pytorch/mae_timm.py                      | 2 ++
 lightly/models/modules/masked_autoencoder_timm.py | 5 ++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/examples/pytorch/mae_timm.py b/examples/pytorch/mae_timm.py
index 00a4ab69e..248675912 100644
--- a/examples/pytorch/mae_timm.py
+++ b/examples/pytorch/mae_timm.py
@@ -1,6 +1,7 @@
 # Note: The model and training settings do not follow the reference settings
 # from the paper. The settings are chosen such that the example can easily be
 # run on a small dataset with a single GPU.
+import sys
 
 import torch
 import torchvision
@@ -14,6 +15,7 @@
     from timm.models import vision_transformer
 except ImportError:
     print("TIMM is not available. Please install in order to run this example.")
+    sys.exit(1)
 
 
 class MAE(nn.Module):
diff --git a/lightly/models/modules/masked_autoencoder_timm.py b/lightly/models/modules/masked_autoencoder_timm.py
index 81205c6f3..238386151 100644
--- a/lightly/models/modules/masked_autoencoder_timm.py
+++ b/lightly/models/modules/masked_autoencoder_timm.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import math
+import sys
 from functools import partial
 from typing import Callable, Optional, Tuple, Type, Union
 
@@ -9,14 +10,12 @@
 except ImportError:
     from typing_extensions import Literal
 
-from timm.layers import LayerType, Mlp, PatchEmbed
-from timm.models import vision_transformer
-
 try:
     from timm.layers import LayerType, Mlp, PatchEmbed
     from timm.models import vision_transformer
 except ImportError:
     print("TIMM is not available. Please install if you would like to use the MAE.")
+    sys.exit(1)
 
 import torch
 import torch.nn as nn