Fix incompatibility between latest skl2onnx and txtai, closes #729

neuml · May 31, 2024 · 5221f73 · 5221f73
1 parent 12fa4f9
commit 5221f73
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/python/txtai/pipeline/train/mlonnx.py b/src/python/txtai/pipeline/train/mlonnx.py
@@ -54,8 +54,9 @@ def __call__(self, model, task="default", output=None, opset=12):
 
         # Find probabilities output node and rename to logits
         for node in model.graph.node:
-            if node.output[0] == "probabilities":
-                node.output[0] = "logits"
+            for x, _ in enumerate(node.output):
+                if node.output[x] == "probabilities":
+                    node.output[x] = "logits"
 
         # Save model to specified output path or return bytes
         model = save_onnx_model(model, output)