Review

Signed-off-by: zethson <[email protected]>

src/scportrait/io/daskmmap.py

@@ -209,4 +209,5 @@ def load_hdf5_chunk(file_path: str, container_name: str, slices: tuple[slice, ..
         The sliced chunk from the HDF5 dataset
     """
     with h5py.File(file_path, "r") as f:
-        f[container_name][slices]
+        data = f[container_name][slices]
+    return data

src/scportrait/pipeline/_utils/spatialdata_classes.py

@@ -104,7 +104,6 @@ def _(self, data: DataTree, classes: set[int] | None = None) -> DataTree:
         if classes is not None:
             for d in data:
                 data[d].attrs["cell_ids"] = classes
-        else:
-            for d in data:
-                data[d] = self._get_cell_ids(data[d])
+        for d in data:
+            data[d] = self._get_cell_ids(data[d])
         return data

src/scportrait/tools/ml/pretrained_models.py

@@ -14,7 +14,7 @@ def load_multilabelSupervised(
     hparam_path: str | Path,
     model_type: str,
     eval: bool = True,
-    device: Literal["cuda", "cpu"] = "cuda",
+    device: Literal["cpu", "cuda", "mps"] = "cuda",
 ) -> MultilabelSupervisedModel:
     """Load a pretrained model uploaded to the github repository.