How can I get the trainer to log steps in dictionary?

Beginners
1

I am training the GPT2 causal LLM huggingface model and the default logging dictionary for each logging step looks like:

{'loss': 9.366, 'grad_norm': 8.703, 'learning_rate': 1e-06, 'epoch': 0.4}

But would like to include the step in this dictionary such as:

{'loss': 9.366, 'grad_norm': 8.703, 'learning_rate': 1e-06, 'epoch': 0.4, 'step': 500}

I have searched the docs and found how I can add metrics like from evaluate, but nothing on how to print predefined things in the training like step. Also, I have not been able to find where the default dictionary values are defined, they don’t seem to be defined in the callback source

I realize you can infer the steps from ‘epoch’, but I would still like to understand so that next time I need to add something to the log I could use this example as a starting point. Thanks!