Bring back necessary tolist call.

Quantco · Feb 13, 2025 · 306a792 · 306a792
1 parent dd9c07e
commit 306a792
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/metalearners/_utils.py b/metalearners/_utils.py
@@ -44,7 +44,7 @@ def index_matrix(matrix: Matrix, rows: Vector) -> Matrix:
     if is_into_dataframe(matrix):
         matrix_nw = nw.from_native(matrix)  # type: ignore
         if rows.dtype == "bool":
-            return matrix_nw[rows].to_native()
+            return matrix_nw[rows.tolist()].to_native()
         return matrix_nw[rows.tolist(), :].to_native()
     return matrix[rows, :]