From b31edf37bffdc186e878ee0c358a28b224db9192 Mon Sep 17 00:00:00 2001 From: Tullie Murrell Date: Fri, 3 Apr 2020 20:59:27 +0200 Subject: [PATCH] Change min max gpu memory to be on their own plots (#1358) --- pytorch_lightning/core/memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/core/memory.py b/pytorch_lightning/core/memory.py index 403e98a094..da31573332 100644 --- a/pytorch_lightning/core/memory.py +++ b/pytorch_lightning/core/memory.py @@ -246,7 +246,7 @@ def get_memory_profile(mode: str) -> Union[Dict[str, int], Dict[int, int]]: min_index, min_memory = min(memory_map.items(), key=lambda item: item[1]) max_index, max_memory = max(memory_map.items(), key=lambda item: item[1]) - memory_map = {min_index: min_memory, max_index: max_memory} + memory_map = {'min_gpu_mem': min_memory, 'max_gpu_mem': max_memory} return memory_map