Update cell_emb.py

cantinilab · Oct 18, 2024 · e30dee5 · e30dee5
1 parent 01e7885
commit e30dee5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scprint/tasks/cell_emb.py b/scprint/tasks/cell_emb.py
@@ -75,6 +75,7 @@ def __init__(
         self.plot_corr_size = plot_corr_size
         self.precision = precision
         self.doplot = doplot
+        self.dtype = dtype
         self.doclass = doclass
         self.output_expression = output_expression
         # subset_hvg=1000, use_layer='counts', is_symbol=True,force_preprocess=True, skip_validate=True)
@@ -153,7 +154,7 @@ def __call__(self, model: torch.nn.Module, adata: AnnData, cache=False):
             model.doplot = self.doplot
             with (
                 torch.no_grad(),
-                torch.autocast(device_type=device, dtype=torch.float16),
+                torch.autocast(device_type=device, dtype=self.dtype),
             ):
                 for batch in tqdm(dataloader):
                     gene_pos, expression, depth = (