Can training step return dictionary?

[igoldov]

My training_step returns {‘loss’: loss, ‘num’: len(batch)}, so I could calculate mean loss at the end of epoch. But now I get this warning, that my training_step returns None and loss doesn't display in the progress bar. How can I return dict from training_step and display loss in the progress bar at the same tim?

[tshu-w]

Can training step return dictionary?

Yes, check the doc of LightningModule.training_step.

How can I return dict from training_step and display loss in the progress bar at the same time?

just return dict like what you do now and use self.log(loss, prog_bar=True) to display loss in prog_bar.

so I could calculate mean loss at the end of epoch

BTW, If you just want mean loss, I think pytorch_lightning have already done for you and you don't need to modified training_step_end(not sure)