trax.supervised.training.Loop.run(n_steps) API

Hello, I have a question related to the API in the title. Link here.

I am confused by how does the n_steps related to batch_size or epochs. Batch_size is specified earlier in the data_generator, epoch is not specified anywhere. Does n_steps mean how many epochs it will go through the entire training set with a certain batch_size? Or it means each step is the size of batch_size, and it will perform n_steps gradient update? I think it is the latter, but in the notebook example, a n_step of 5 will not even get close to finish one epoch. Am I correct? Thanks.

Hi @Longyu_Zhao

Yes you are correct - steps mean weight updates (each batch).