fix bug

MannLabs · Nov 29, 2024 · 60a6818 · 60a6818
1 parent 9bca6e0
commit 60a6818
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/scportrait/pipeline/classification.py b/src/scportrait/pipeline/classification.py
@@ -403,7 +403,7 @@ def inference(self,
 
             #save the results for each batch into the memory mapped array at the specified indices
             features[ix:(ix+batch_size)] = result.numpy()
-            cell_ids[ix:(ix+batch_size)] = class_id.numpy().unsqueeze(1)
+            cell_ids[ix:(ix+batch_size)] = class_id.unsqueeze(1)
 
             for i in range(len(dataloader) - 1):
                 if i % 10 == 0:
@@ -415,7 +415,7 @@ def inference(self,
 
                 #save the results for each batch into the memory mapped array at the specified indices
                 features[ix:(ix+batch_size)] = r.cpu().detach().numpy()
-                cell_ids[ix:(ix+batch_size)] = label.numpy().unsqueeze(1)
+                cell_ids[ix:(ix+batch_size)] = label.unsqueeze(1)
 
         if hasattr(self.config, "log_transform"):
             if self.config["log_transform"]: