fix a bug

lucidrains · Jul 25, 2024 · 6d797df · 6d797df
1 parent f24863f
commit 6d797df
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/autoregressive_diffusion_pytorch/autoregressive_diffusion.py b/autoregressive_diffusion_pytorch/autoregressive_diffusion.py
@@ -464,9 +464,10 @@ def sample(
         for _ in tqdm(range(self.max_seq_len - out.shape[1]), desc = 'tokens'):
 
             cond = self.proj_in(out)
-            cond = cond + self.abs_pos_emb(torch.arange(cond.shape[1], device = self.device))
 
             cond = torch.cat((start_tokens, cond), dim = 1)
+            cond = cond + self.abs_pos_emb(torch.arange(cond.shape[1], device = self.device))
+
             cond = self.transformer(cond)
 
             last_cond = cond[:, -1]
@@ -494,8 +495,8 @@ def forward(
 
         seq = self.proj_in(seq)
         start_token = repeat(self.start_token, 'd -> b 1 d', b = b)
-        seq = torch.cat((start_token, seq), dim = 1)
 
+        seq = torch.cat((start_token, seq), dim = 1)
         seq = seq + self.abs_pos_emb(torch.arange(seq_len, device = self.device))
 
         cond = self.transformer(seq)

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "autoregressive-diffusion-pytorch"
-version = "0.0.6"
+version = "0.1.0"
 description = "Autoregressive Diffusion - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }