Make sure saved `onnx` models support batching out of the box

When loading a `onnx` model after training via the standard pipeline, it seems that batching is not supported out of the box (no leading batch dimension).

Need to investigate this.

Start by tracing the standard training pipeline through saving and loading the ONNX model, focusing on where input dimensions are defined. Reproduce the issue with a batched input; done means a saved model loaded through the standard pipeline accepts a leading batch dimension out of the box.

Written by the indexing model from the issue text.