Add batching in order to allow for different fit and predict inputs

[monocongo]

We currently expect data for prediction to match the shape of the fit inputs, i.e. there is currently no allowance for batching in the model(s).