final cleanup

lucidrains · Jan 19, 2021 · 1162df5 · 1162df5
1 parent 8960885
commit 1162df5
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/dalle_pytorch/dalle_pytorch.py b/dalle_pytorch/dalle_pytorch.py
@@ -374,7 +374,7 @@ def forward(
         mask = None,
         return_loss = False
     ):
-        device = text.device
+        device, ignore_index = text.device, self.ignore_index
 
         text = F.pad(text, (1, 0), value = 0) # use padding as <bos>
         tokens = self.text_emb(text)
@@ -401,9 +401,9 @@ def forward(
         logits = self.to_logits(out)
 
         # mask logits to make sure text predicts text (except last token), and image predicts image
-        mask = self.logits_mask[:, :seq_len]
+        logits_mask = self.logits_mask[:, :seq_len]
         max_neg_value = -torch.finfo(logits.dtype).max
-        logits.masked_fill_(mask, max_neg_value)
+        logits.masked_fill_(logits_mask, max_neg_value)
 
         if not return_loss:
             return logits
@@ -415,8 +415,9 @@ def forward(
 
         if noncausal_attn_len > 0:
             seq_range = torch.arange(seq_len, device = device)
-            mask = seq_range < noncausal_attn_len
-            labels.masked_fill_(mask[None, :], -100) # -100 is the ignore index for cross entropy loss
+            noncausal_attn_mask = seq_range < noncausal_attn_len
+            noncausal_attn_mask = rearrange(noncausal_attn_mask, 'n -> () n')
+            labels.masked_fill_(noncausal_attn_mask, ignore_index)
 
         loss = F.cross_entropy(rearrange(logits, 'b n c -> b c n'), labels)
         return loss
diff --git a/dalle_pytorch/transformer.py b/dalle_pytorch/transformer.py
@@ -87,7 +87,7 @@ def forward(self, x, mask = None):
             i, j = dots.shape[-2:]
             mask = torch.ones(i, j, device = device).triu_(j - i + 1).bool()
 
-            if self.noncausal_attn_len > 0:
+            if self.noncausal_attn_len > 1:
                 ind = slice(0, self.noncausal_attn_len)
                 mask[ind, ind] = False
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'dalle-pytorch',
   packages = find_packages(),
-  version = '0.0.45',
+  version = '0.0.46',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',