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

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