The code for data and training components has been refactored using Lightning to support multi-GPU parallelism