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.