make it work for @ethancohen123

Author: lucidrains

dalle2_pytorch/dalle2_pytorch.py

@@ -250,9 +250,13 @@ def embed_text(self, text):
         text = text[..., :self.max_text_len]
         text_mask = text != 0
         encoder_output = self.clip.text_transformer(text)
-        text_cls, text_encodings = encoder_output[:, 0], encoder_output[:, 1:]
+
+        text_cls, text_encodings = (encoder_output[:, 0], encoder_output[:, 1:]) if encoder_output.ndim == 3 else (encoder_output, None)
         text_embed = self.clip.to_text_latent(text_cls)
-        text_encodings = text_encodings.masked_fill(~text_mask[..., None], 0.)
+
+        if exists(text_encodings):
+            text_encodings = text_encodings.masked_fill(~text_mask[..., None], 0.)
+
         return EmbeddedText(l2norm(text_embed), text_encodings)
 
     @torch.no_grad()

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.2'
+__version__ = '1.8.3'