Add batched prediction to molecular dynamics

[miketynes]

• add a batch_size parameter to the MolecularDynamics class
• pass it to the Predictor.__call__ method during each _step of molecular dynamics

The motivation for this is to prevent OOM errors encountered when doing large ensemble sampling runs.

[lubbersnick]

Looks good, just going to hold it open a bit in case you find something very similar and relevant for us to bundle in.