Training of a universal differential equation on multiple initial states and mini-batches

[xhemalkodragjini]

Hi,
I am training a UDE controlled by an input signal as part of my thesis. To be able to get a more generalized neural network on multiple data, I want to train on around 1000 different initial states and use mini-batches on the go.
I was not able to find some example of training with multiple initial states so far. Is there any available? If not, what approach should be used?

The UDE contains some biological laws which are known and some unknowns which will be approximated by the NN.

Many thanks in advance