From 5221f732277d39d396c6cbdcc323486720e286de Mon Sep 17 00:00:00 2001 From: davidmezzetti <561939+davidmezzetti@users.noreply.github.com> Date: Fri, 31 May 2024 08:03:55 -0400 Subject: [PATCH] Fix incompatibility between latest skl2onnx and txtai, closes #729 --- src/python/txtai/pipeline/train/mlonnx.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/python/txtai/pipeline/train/mlonnx.py b/src/python/txtai/pipeline/train/mlonnx.py index 1e3bac653..126e025c7 100644 --- 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)