Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

This commit is contained in:
Tullie Murrell 2020-04-03 20:59:27 +02:00 committed by GitHub
parent 003ca510fa
commit b31edf37bf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pytorch_lightning/core/memory.py
View File

@ -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