Fixed bug with array conversion

MinishLab · Feb 15, 2025 · 5c9d397 · 5c9d397
1 parent ff4043f
commit 5c9d397
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model2vec/inference/model.py b/model2vec/inference/model.py
@@ -134,7 +134,7 @@ def predict(
             proba = self.head.predict_proba(encoded)
             for vector in proba:
                 out_labels.append(self.classes_[vector > threshold])
-            return np.asarray(out_labels)
+            return np.asarray(out_labels, dtype=object)
 
         return self.head.predict(encoded)