take care of receiving single dna sequence

lucidrains · Oct 28, 2021 · 230e6c9 · 230e6c9
1 parent 074e2f9
commit 230e6c9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/enformer_pytorch/enformer_pytorch.py b/enformer_pytorch/enformer_pytorch.py
@@ -333,8 +333,17 @@ def heads(self):
         return self._heads
 
     def forward(self, x):
+        no_batch = (x.ndim == 1)
+        if no_batch:
+            x = rearrange(x, 'n -> () n')
+
         if x.dtype == torch.long:
             x = F.one_hot(x, num_classes = self.num_alphabet)
 
         x = self._trunk(x.float())
-        return map_values(lambda fn: fn(x), self._heads)
+        out = map_values(lambda fn: fn(x), self._heads)
+
+        if no_batch:
+            out = map_values(lambda t: rearrange(t, '() ... -> ...'), out)
+
+        return out
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.7',
+  version = '0.0.8',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',