fixed import

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):

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