Can ONNX Runtime Training Handle Models with Different Outputs for Training and Inference?

[Leo5050xvjf]

Hello,

If my model's structure differs slightly between training and inference phases, is it still suitable to use ONNX Runtime Training? For instance, during training, my model has two outputs ( D_q and D_q_), but during inference, it only has one output (D_q).

In this case, can I finish training on the device, export the inference model, and then remove D_q_ and any unused nodes from the inference model?

Thanks!

Here is a sample image.

(Model with two output)

(Model Only One output)