Based on the docs of torch.onnx.export it seems you have the following choices for the training argument:

training (enum , default TrainingMode.EVAL) –


TrainingMode.EVAL: export the model in inference mode.

TrainingMode.PRESERVE: export the model in inference mode if model.training is False and in training mode if model.training is True.

TrainingMode.TRAINING: export the model in training mode. Disables optimizations which might interfere with training.


Based on this I don’t think the exported model can switch between training and evaluation and you would need to pick one option.