new einops automatically torch compile friendly

lucidrains · Oct 18, 2023 · 33d2aad · 33d2aad
1 parent 4a34864
commit 33d2aad
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 9 deletions.
diff --git a/imagen_pytorch/__init__.py b/imagen_pytorch/__init__.py
@@ -1,10 +1,3 @@
-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.10'
+__version__ = '1.25.11'
diff --git a/setup.py b/setup.py
@@ -30,7 +30,7 @@
     'beartype',
     'click',
     'datasets',
-    'einops>=0.6.1',
+    'einops>=0.7.0',
     'ema-pytorch>=0.0.3',
     'fsspec',
     'kornia',