fix get_at usage

lucidrains · Jan 20, 2024 · 3d5a1b5 · 3d5a1b5
1 parent f6d36da
commit 3d5a1b5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/equiformer_pytorch/equiformer_pytorch.py b/equiformer_pytorch/equiformer_pytorch.py
@@ -337,9 +337,7 @@ def forward(
 
                 xi, xj = source[degree_in], target[degree_in]
 
-                flattened_neighbor_indices, ps = pack_one(neighbor_indices, 'b *')
-                x = get_at('b [i] d m, b k -> b k d m', xj, flattened_neighbor_indices)
-                x = unpack_one(x, ps, 'b * d m')
+                x = get_at('b [i] d m, b j k -> b j k d m', xj, neighbor_indices)
 
                 if self.project_xi_xj:
                     xi = rearrange(xi, 'b i d m -> b i 1 d m')

diff --git a/equiformer_pytorch/version.py b/equiformer_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.5.2'
+__version__ = '0.5.3'