Issue fixed and vectorized some parts of train_cnn