address #21

Author: lucidrains

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

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',