Improve DataModule flexibility #752
Description
As in the current version, the DataModule is not flexible enough. Specifically, the user cannot specify how batches are composed and how to retrive data from conditions. Enhancing DataModule will allow to use PINA in many other context such as Autoregressive training (see #730)