Add training and inference code

Add train and custom dataloader

initial commit