Updated type check

MinishLab · Feb 16, 2025 · f9037d9 · f9037d9
1 parent eccec80
commit f9037d9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py
@@ -231,9 +231,9 @@ def _initialize(self, y: LabelType) -> None:
         :raises ValueError: If the labels are inconsistent.
         """
         if isinstance(y[0], (str, int)):
-            # Check if all labels are strings.
+            # Check if all labels are strings or integers.
             if not all(isinstance(label, (str | int)) for label in y):
-                raise ValueError("Inconsistent label types in y. All labels must be strings.")
+                raise ValueError("Inconsistent label types in y. All labels must be strings or integers.")
             self.multilabel = False
             classes = sorted(set(y))
         else: