address lucidrains#21

ewonglab · Feb 11, 2023 · 4e70710 · 4e70710
1 parent abb50a0
commit 4e70710
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/enformer_pytorch/modeling_enformer.py b/enformer_pytorch/modeling_enformer.py
@@ -136,8 +136,14 @@ def __init__(self, dim, pool_size = 2):
         super().__init__()
         self.pool_size = pool_size
         self.pool_fn = Rearrange('b d (n p) -> b d n p', p = pool_size)
+
         self.to_attn_logits = nn.Conv2d(dim, dim, 1, bias = False)
 
+        nn.init.dirac_(self.to_attn_logits.weight)
+
+        with torch.no_grad():
+            self.to_attn_logits.weight.mul_(2)
+
     def forward(self, x):
         b, _, n = x.shape
         remainder = n % self.pool_size

diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.6.2',
+  version = '0.6.4',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',