Train data function made and added wandb as a dependency

[lsig]

Train data function made and type annotated

Made the train data function inside the training folder for training the model.

How to test

• Ran format_and_lint before pushing
• No unit tests