Function batched_inference 

pub fn batched_inference<F>(
    input: &Tensor,
    config: &DynamicBatchConfig,
    infer_fn: F,
) -> Result<Tensor, ModelError>
where
    F: Fn(&Tensor) -> Result<Tensor, ModelError>,

Splits a large input into batches, runs inference, and reassembles.

input is [N, ...], infer_fn processes [B, ...] and returns [B, ...].