Fix doctests.

aai-institute · Jul 25, 2023 · 3720cd7 · 3720cd7
1 parent 495211e
commit 3720cd7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/pydvl/utils/parallel/backend.py b/src/pydvl/utils/parallel/backend.py
@@ -212,7 +212,8 @@ def init_parallel_backend(
     >>> config = ParallelConfig(backend="ray")
     >>> parallel_backend = init_parallel_backend(config)
     >>> parallel_backend
-    <RayParallelBackend: {'address': None, 'logging_level': 30, 'num_cpus': None}>
+    <RayParallelBackend: {'address': None, 'logging_level': 30, '_temp_dir': None, \
+'num_cpus': None}>
 
     """
     try:

diff --git a/src/pydvl/utils/utility.py b/src/pydvl/utils/utility.py
@@ -189,7 +189,11 @@ def _utility(self, indices: FrozenSet) -> float:
                     model = self.model
 
                 # Special case for classification with only one class
-                if y_train.dtype == int and len(np.unique(y_train)) == 1:
+                if (
+                    y_train.dtype == int
+                    and len(np.unique(y_train.reshape(-1))) == 1
+                    and self.scorer._name == "accuracy"
+                ):
                     unique_cls = y_train[0]
                     all_cls, counts = np.unique(y_test, return_counts=True)
                     cls_idx = np.argwhere(all_cls == unique_cls)[0, 0]