Batching predictions on targets

Is there any way to simultaneously get predictions for multiple timesteps from the Graphcast model? We have some training infrastructure set up and batching in place, but we need a way to execute the model on multiple inputs to be able to train it.