This milestone covers all tasks related to refactoring the dataset handling to align with PyTorch Lightning standards. The goal is to ensure that data pipelines are modular, scalable, and easy to integrate into the PyTorch Lightning workflow.