Switch timm import fallback order to favour new paths, avoid new warn…

…ings
mlfoundations · Oct 17, 2024 · 236585c · 236585c
1 parent 6de2025
commit 236585c
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/src/open_clip/timm_model.py b/src/open_clip/timm_model.py
@@ -10,15 +10,16 @@
 
 try:
     import timm
-    from timm.models.layers import Mlp, to_2tuple
     try:
-        # old timm imports < 0.8.1
-        from timm.models.layers.attention_pool2d import RotAttentionPool2d
-        from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
-    except ImportError:
         # new timm imports >= 0.8.1
         from timm.layers import RotAttentionPool2d
         from timm.layers import AttentionPool2d as AbsAttentionPool2d
+        from timm.layers import Mlp, to_2tuple
+    except ImportError as e:
+        # fallback, try old timm imports < 0.8.1
+        from timm.models.layers.attention_pool2d import RotAttentionPool2d
+        from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
+        from timm.models.layers import Mlp, to_2tuple
 except ImportError:
     timm = None