From dccab6ce19f03f7b0b951c82f271e3e0cbfed22d Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sat, 5 Oct 2019 16:38:48 -0400 Subject: [PATCH] cleaning up demos --- pytorch_lightning/trainer/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 474080f0..f79cdb1a 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -1132,7 +1132,7 @@ class Trainer(TrainerIO): # add gpu memory if self.on_gpu and self.log_gpu_memory: - mem_map = memory.get_memory_profile() + mem_map = memory.get_memory_profile(self.log_gpu_memory) metrics.update(mem_map) # add norms