fix ci

huggingface · Feb 12, 2025 · 23f6c9b · 23f6c9b
1 parent 893ca6b
commit 23f6c9b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
@@ -1340,8 +1340,8 @@ def _add_mask(
 ):
     if isinstance(input, pa.Table):
         if not isinstance(mask, (list, pa.Array, pa.ChunkedArray)):
-            mask = [mask]
-        return input.add_column(mask_column_name, mask)
+            mask = pa.array([mask], type=pa.bool_())
+        return input.append_column(mask_column_name, mask)
     else:
         return {mask_column_name: mask}