make sure einops is 0.6.1 or greater and call allow_ops_in_compiled_g…

…raph
lucidrains · Apr 20, 2023 · 426eb34 · 426eb34
1 parent 2edd8aa
commit 426eb34
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/imagen_pytorch/__init__.py b/imagen_pytorch/__init__.py
@@ -1,3 +1,10 @@
+import torch
+from packaging import version
+
+if version.parse(torch.__version__) >= version.parse('2.0.0'):
+    from einops._torch_specific import allow_ops_in_compiled_graph
+    allow_ops_in_compiled_graph()
+
 from imagen_pytorch.imagen_pytorch import Imagen, Unet
 from imagen_pytorch.imagen_pytorch import NullUnet
 from imagen_pytorch.imagen_pytorch import BaseUnet64, SRUnet256, SRUnet1024

diff --git a/imagen_pytorch/version.py b/imagen_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.25.0'
+__version__ = '1.25.2'
diff --git a/setup.py b/setup.py
@@ -30,7 +30,7 @@
     'beartype',
     'click',
     'datasets',
-    'einops>=0.6',    
+    'einops>=0.6.1',
     'ema-pytorch>=0.0.3',
     'fsspec',
     'kornia',