Update lightning_module_template.py

2019-10-23 04:31:58 -04:00 · 2019-10-23 04:31:58 -04:00 · 3a2466258d
parent c5c03c87db
commit 3a2466258d
1 changed files with 1 additions and 1 deletions
--- a/pl_examples/basic_examples/lightning_module_template.py
+++ b/pl_examples/basic_examples/lightning_module_template.py
@ -158,7 +158,7 @@ class LightningTemplateModel(LightningModule):
            val_loss = output['val_loss']

            # reduce manually when using dp
-            if self.trainer.use_dp:
+            if self.trainer.use_dp or self.trainer.use_ddp2:
                val_loss = torch.mean(val_loss)
            val_loss_mean += val_loss