Implement Dataclasses

• Read through the code
• Implement dataclasses instead of vectors and matrices
  • Within dataset_iter_multi.py for better storing the data
  • Within trainer.py to use it inside the training loop
• update collate fn
• use dataclasses within the ae architecture
• modify the loss fn to process dataclasses, not jax arrays