go with what @cfoster0 says

lucidrains · Mar 5, 2021 · 005279f · 005279f
1 parent cf4d8d9
commit 005279f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/glom_pytorch/glom_pytorch.py b/glom_pytorch/glom_pytorch.py
@@ -93,7 +93,9 @@ def forward(self, img, iters = None, return_all = False):
         n = tokens.shape[1]
 
         pos_embs = self.pos_emb(torch.arange(n, device = device))
-        bottom_level = tokens + rearrange(pos_embs, 'n d -> () n d')
+        pos_embs = rearrange(pos_embs, 'n d -> () n () d')
+
+        bottom_level = tokens
         bottom_level = rearrange(bottom_level, 'b n d -> b n () d')
 
         levels = repeat(self.init_levels, 'l d -> b n l d', b = b, n = n)
@@ -104,7 +106,7 @@ def forward(self, img, iters = None, return_all = False):
             levels_with_input = torch.cat((bottom_level, levels), dim = -2)  # each iteration, attach original input (with positional embedding) at the bottom level
 
             bottom_up_out = self.bottom_up(levels_with_input[..., 1:-1, :])
-            top_down_out = self.top_down(levels_with_input[..., 2:, :])
+            top_down_out = self.top_down(levels_with_input[..., 2:, :] + pos_embs) # positional embeddings given to top-down networks
 
             bottom_up_out = torch.cat((bottom_level, bottom_up_out), dim = -2)
             top_down_out = F.pad(top_down_out, (0, 0, 0, 1), value = 0.)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'glom-pytorch',
   packages = find_packages(),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'Glom - Pytorch',
   author = 'Phil Wang',