In-memory trajectory dataset and custom batch creation

[andrrizzi]

MDAnalysis loads each frame from the disk. For big trajectories, it might be worth it support in-memory trajectories (loaded after subsampling and atom selections). MDAnalysis also have a utility to load the trajectory in memory that is worth checking out before implementing our logic.

An advantage of this implementation is also that a custom DataLoader might support creating batches by slicing the tensors rather than accessing the samples one by one and concatenating them.