minor changes

feat(model, train): improved architecture, overfit prevention, re-trained the model

feat(data): implemented data augmentation

feat(train, dist_train, plots): improved training, implemented distributed training, added plotting