added training v2 + few fixes to training script