Updated type check

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