add an assert on torch version in order to use OpenAIDiscreteVAE

dalle_pytorch/vae.py

@@ -10,6 +10,8 @@
 from pathlib import Path
 from tqdm import tqdm
 from math import sqrt, log
+from packaging import version
+
 from omegaconf import OmegaConf
 from taming.models.vqgan import VQModel, GumbelVQ
 import importlib
@@ -98,11 +100,18 @@ def make_contiguous(module):
         for param in module.parameters():
             param.set_(param.contiguous())
 
+# package versions
+
+def get_pkg_version(pkg_name):
+    from pkg_resources import get_distribution
+    return get_distribution(pkg_name).version
+
 # pretrained Discrete VAE from OpenAI
 
 class OpenAIDiscreteVAE(nn.Module):
     def __init__(self):
         super().__init__()
+        assert version.parse(get_pkg_version('torch')) < version.parse('1.11.0'), 'torch version must be <= 1.10 in order to use OpenAI discrete vae'
 
         self.enc = load_model(download(OPENAI_VAE_ENCODER_PATH))
         self.dec = load_model(download(OPENAI_VAE_DECODER_PATH))

setup.py

@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '1.6.1',
+  version = '1.6.2',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',