fall in line with trim logic in enformer master

lucidrains · Jun 1, 2022 · f881928 · f881928
1 parent dcb5751
commit f881928
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/enformer_pytorch/modeling_enformer.py b/enformer_pytorch/modeling_enformer.py
@@ -174,6 +174,10 @@ def forward(self, x):
             raise ValueError(f'sequence length {seq_len} is less than target length {target_len}')
 
         trim = (target_len - seq_len) // 2
+
+        if trim == 0:
+            return x
+
         return x[:, -trim:trim]
 
 def ConvBlock(dim, dim_out = None, kernel_size = 1):

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