Is there any way to save loss of training step?

[Derick317]

I wonder there is any way to save loss of training step? Since I install the repository from pip, changing the souce code seems dirty.

[kirilzilla]

in class Trainer, def train() you can for example import wandb, init and after

pbar.set_description(f'loss: {total_loss:.4f}')

Log the total_loss to wandb

wandb.log({"total_loss": total_loss}, step=self.step)