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

Lightning-AI · Apr 3, 2020 · b31edf3 · b31edf3
1 parent 003ca51
commit b31edf3
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