make sure generate_text works when shift token feature is turned on

Author: lucidrains

dalle_pytorch/transformer.py

@@ -103,6 +103,11 @@ def forward(self, x, **kwargs):
         text_len = seq_len - img_seq_len + 1
         padding = seq_len - n + 1
 
+        # if sequence is shorter than the text length, no image tokens to shift
+
+        if n < text_len:
+            return self.fn(x, **kwargs)
+
         # get text and image tokens
 
         x_text, x_img = x[:, :text_len], x[:, text_len:]

setup.py

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