Change min max gpu memory to be on their own plots (Lightning-AI#1358)

alexeykarnachev · Apr 4, 2020 · 1f10248 · 1f10248
1 parent c2c88f5
commit 1f10248
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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