From 170262ff738d7d23365ea5f7c94b832009992b5b Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sat, 5 Oct 2019 15:54:48 -0400 Subject: [PATCH] cleaning up demos --- examples/basic_examples/lightning_module_template.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/basic_examples/lightning_module_template.py b/examples/basic_examples/lightning_module_template.py index 5cb2c3ee..88b8534c 100644 --- a/examples/basic_examples/lightning_module_template.py +++ b/examples/basic_examples/lightning_module_template.py @@ -95,7 +95,7 @@ class LightningTemplateModel(LightningModule): loss_val = self.loss(y, y_hat) # in DP mode (default) make sure if result is scalar, there's another dim in the beginning - if self.trainer.use_dp: + if self.trainer.use_dp or self.trainer.use_ddp2: loss_val = loss_val.unsqueeze(0) output = OrderedDict({ @@ -126,7 +126,7 @@ class LightningTemplateModel(LightningModule): val_acc = val_acc.cuda(loss_val.device.index) # in DP mode (default) make sure if result is scalar, there's another dim in the beginning - if self.trainer.use_dp: + if self.trainer.use_dp or self.trainer.use_ddp2: loss_val = loss_val.unsqueeze(0) val_acc = val_acc.unsqueeze(0) @@ -168,7 +168,7 @@ class LightningTemplateModel(LightningModule): val_loss_mean /= len(outputs) val_acc_mean /= len(outputs) tqdm_dict = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean} - result = {'progress_bar': tqdm_dict} + result = {'progress_bar': tqdm_dict, 'logs': tqdm_dict} return result # ---------------------