Update lightning_module_template.py

2026-09-09 11:32:07 +08:00 · 2019-10-07 17:23:25 -04:00
parent c0b0c91d24
commit 46b55d9aaa
1 changed files with 4 additions and 1 deletions
@@ -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)