fix bug

lucidrains · Jun 24, 2022 · e7a31c2 · e7a31c2
1 parent b301b74
commit e7a31c2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/parti_pytorch/parti_pytorch.py b/parti_pytorch/parti_pytorch.py
@@ -301,7 +301,7 @@ def generate(
 
         image_tokens = torch.empty((batch, 0), device = device, dtype = torch.long)
 
-        for _ in range(image_seq_len + 1):
+        for _ in range(image_seq_len):
             logits = self.forward_with_cond_scale(
                 text_token_embeds = text_token_embeds,
                 text_mask = text_mask,
@@ -314,8 +314,6 @@ def generate(
             sampled = rearrange(sampled, 'b -> b 1')
             image_tokens = torch.cat((image_tokens, sampled), dim = -1)
 
-        image_tokens = image_tokens[:, 1:] # remove start token
-
         image_tokens = rearrange(image_tokens, 'b (h w) -> b h w', h = self.image_encoded_dim)
 
         if not exists(self.vae):

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'parti-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'Parti - Pathways Autoregressive Text-to-Image Model - Pytorch',
   author = 'Phil Wang',