Merge pull request #3 from lucidrains/modulate-feature-dimension

different attention map for each element of feature dimension
lucidrains · Jan 18, 2021 · b8fcbc2 · b8fcbc2
2 parents 7860b28 + a0a56d8
commit b8fcbc2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/point_transformer_pytorch/point_transformer_pytorch.py b/point_transformer_pytorch/point_transformer_pytorch.py
@@ -23,7 +23,7 @@ def __init__(
         self.attn_mlp = nn.Sequential(
             nn.Linear(dim, dim * attn_mlp_hidden_mult),
             nn.ReLU(),
-            nn.Linear(dim * attn_mlp_hidden_mult, 1),
+            nn.Linear(dim * attn_mlp_hidden_mult, dim),
         )
 
     def forward(self, x, pos):
@@ -47,8 +47,8 @@ def forward(self, x, pos):
         v = v + rel_pos_emb
 
         # attention
-        attn = sim.softmax(dim = -1)
+        attn = sim.softmax(dim = -2)
 
         # aggregate
-        agg = einsum('b i j, b i j d -> b i d', attn, v)
+        agg = einsum('b i j d, b i j d -> b i d', attn, v)
         return agg
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'point-transformer-pytorch',
   packages = find_packages(),
-  version = '0.0.1',
+  version = '0.0.2',
   license='MIT',
   description = 'Point Transformer - Pytorch',
   author = 'Phil Wang',