Abstract Inspired by the iterative procedure of computing mean fields with known Reynolds stresses (Theoret. Appl. Mech. Lett., 2021), we proposed a way to achieve data augmentation by utilizing the intermediate mean fields after proper selections. We also proposed modifications to the Tensor Basis Neural Network (J. Fluid Mech., 2016) model. With the modification of the learning targets and the inclusions of wall distance and logarithm of normalized eddy viscosity in the model inputs, the modified version of the model with augmented training datasets shows better performance on Reynolds stress predictions for two dimensional incompressible flow over periodic hills under different geometries. Furthermore, better propagated mean velocity fields can be achieved, showing better agreements with the DNS results.