From 1c45edd192e48829ae2d8c0ea2f8795609ea2908 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sat, 5 Oct 2019 16:30:27 -0400 Subject: [PATCH] cleaning up demos --- pytorch_lightning/root_module/memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py index 6aa4519e..ec4e2c02 100644 --- a/pytorch_lightning/root_module/memory.py +++ b/pytorch_lightning/root_module/memory.py @@ -35,7 +35,7 @@ class ModelSummary(object): out_sizes = [] input_ = self.model.example_input_array - if self.model.use_ddp or self.model.use_dp or self.model.single_gpu: + if self.model.on_gpu: input_ = input_.cuda(0) if self.model.trainer.use_amp: