release improved coordinate normalization, thanks to @MattMcPartlon

lucidrains · Jun 4, 2021 · 44c3687 · 44c3687
1 parent bf2bda8
commit 44c3687
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/egnn_pytorch/egnn_pytorch.py b/egnn_pytorch/egnn_pytorch.py
@@ -88,13 +88,13 @@ class CoorsNorm(nn.Module):
     def __init__(self, eps = 1e-8):
         super().__init__()
         self.eps = eps
-        self.fn = nn.LayerNorm(1)
+        self.scale = nn.Parameter(torch.ones(1))
+        self.bias = nn.Parameter(torch.zeros(1))
 
     def forward(self, coors):
         norm = coors.norm(dim = -1, keepdim = True)
         normed_coors = coors / norm.clamp(min = self.eps)
-        phase = self.fn(norm)
-        return (phase * normed_coors)
+        return normed_coors * self.scale + self.bias
 
 # classes
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'egnn-pytorch',
   packages = find_packages(),
-  version = '0.1.12',
+  version = '0.2.0',
   license='MIT',
   description = 'E(n)-Equivariant Graph Neural Network - Pytorch',
   author = 'Phil Wang, Eric Alcaide',