Fixed merge conflict

MinishLab · Feb 16, 2025 · 5003768 · 5003768
2 parents 2dc5b17 + 2d51516
commit 5003768
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py
@@ -96,7 +96,11 @@ def predict(
                 pred.extend([self.classes[np.flatnonzero(row)] for row in mask])
             else:
                 pred.extend([self.classes[idx] for idx in logits.argmax(dim=1).tolist()])
-        return np.array(pred, dtype=object)
+        if self.multilabel:
+            # Return as object array to allow for lists of varying lengths.
+            return np.array(pred, dtype=object)
+        else:
+            return np.array(pred)
 
     @torch.no_grad()
     def _predict_single_batch(self, X: list[str]) -> torch.Tensor: