convert indices to int8 before stacking

Genentech · Jul 16, 2024 · 7b20023 · 7b20023
1 parent c5f3c14
commit 7b20023
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/grelu/sequence/format.py b/src/grelu/sequence/format.py
@@ -252,8 +252,11 @@ def strings_to_indices(
             strings
         ), "All input sequences must have the same length."
         return np.stack(
-            [[BASE_TO_INDEX_HASH[base] for base in string] for string in strings]
-        ).astype(np.int8)
+            [
+                np.array([BASE_TO_INDEX_HASH[base] for base in string], dtype(np.int8))
+                for string in strings
+            ]
+        )
 
 
 def indices_to_one_hot(indices: np.ndarray) -> Tensor: