Exported ONNX Model doesn't output matched jet index predictions

[jpgombas]

When predicting assignments through spanet.predict, you get for example "/TARGETS/tR/b" in the output h5 file. This contains the index that spanet predicts a jet to match to the b from "tR". However, when exporting, the ONNX model doesn't have any predicted assignments. It does have the detection and assignment probabilities though. Is there a way to get this prediction through the embeddings, or is this feature currently unavailable?