address #21

lucidrains · Feb 11, 2023 · 5dc54a8 · 5dc54a8
1 parent abb50a0
commit 5dc54a8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/enformer_pytorch/modeling_enformer.py b/enformer_pytorch/modeling_enformer.py
@@ -136,7 +136,9 @@ 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)
 
     def forward(self, x):
         b, _, n = x.shape

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.3',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',