From 46b55d9aaa494773301be67af41a1b8bebf55960 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 7 Oct 2019 17:23:25 -0400 Subject: [PATCH] Update lightning_module_template.py --- examples/basic_examples/lightning_module_template.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/basic_examples/lightning_module_template.py b/examples/basic_examples/lightning_module_template.py index 88b8534c..a6d94595 100644 --- a/examples/basic_examples/lightning_module_template.py +++ b/examples/basic_examples/lightning_module_template.py @@ -98,8 +98,11 @@ class LightningTemplateModel(LightningModule): if self.trainer.use_dp or self.trainer.use_ddp2: loss_val = loss_val.unsqueeze(0) + tqdm_dict = {'train_loss': loss_val} output = OrderedDict({ - 'loss': loss_val + 'loss': loss_val, + 'progress_bar': tqdm_dict, + 'loss': tqdm_dict }) # can also return just a scalar instead of a dict (return loss_val)