fix tests

DeplanckeLab · Aug 5, 2023 · 8d6ab6b · 8d6ab6b
1 parent 522e97c
commit 8d6ab6b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/chromatinhd/models/pred/loader/transcriptome.py b/src/chromatinhd/models/pred/loader/transcriptome.py
@@ -17,9 +17,12 @@ class Transcriptome:
     def __init__(
         self,
         transcriptome: chromatinhd.data.transcriptome.Transcriptome,
-        layer: str = "X",
+        layer: str = None,
     ):
-        X = transcriptome.layers[layer]
+        if layer is None:
+            X = transcriptome.X
+        else:
+            X = transcriptome.layers[layer]
         if chromatinhd.sparse.is_sparse(X):
             self.X = X.dense()
         elif torch.is_tensor(X):

diff --git a/src/chromatinhd/models/pred/loader/transcriptome_fragments.py b/src/chromatinhd/models/pred/loader/transcriptome_fragments.py
@@ -25,7 +25,7 @@ def __init__(
         fragments: chromatinhd.data.fragments.Fragments,
         transcriptome: chromatinhd.data.transcriptome.Transcriptome,
         cellxgene_batch_size: int,
-        layer: str = "X",
+        layer: str = None,
     ):
         # ensure that transcriptome and fragments have the same var
         if not all(transcriptome.var.index == fragments.var.index):