make sure last token can still predict eos

QROST · Jan 7, 2021 · c32f724 · c32f724
1 parent 7582af4
commit c32f724
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/dalle_pytorch/dalle_pytorch.py b/dalle_pytorch/dalle_pytorch.py
@@ -256,6 +256,7 @@ def __init__(
  self.text_seq_len = text_seq_len
  self.image_seq_len = image_seq_len
 
+ seq_len = text_seq_len + image_seq_len
  total_tokens = num_text_tokens + num_image_tokens + 1 # extra for EOS
  self.total_tokens = total_tokens
 
@@ -271,7 +272,7 @@ def __init__(
  nn.Linear(dim, self.total_tokens),
  )
 
- seq_range = torch.arange(text_seq_len + image_seq_len)
+ seq_range = torch.arange(seq_len)
  logits_range = torch.arange(total_tokens)
 
  seq_range = rearrange(seq_range, 'n -> () n ()')
@@ -280,7 +281,7 @@ def __init__(
  logits_mask = (
  ((seq_range >= (text_seq_len - 1)) & (logits_range < num_text_tokens)) |
  ((seq_range < (text_seq_len - 1)) & (logits_range >= num_text_tokens)) |
- (logits_range >= (total_tokens - 1))
+ ((seq_range != (seq_len - 1)) & (logits_range >= (total_tokens - 1)))
  )
 
  self.register_buffer('logits_mask', logits_mask)

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