utils/debug_option.py

import torch
DEBUG = False
DEBUG_GRADIENT = False
LOG_GPU_MEMORY = False
_GPU_MEM_GB = 0.0
if torch.cuda.is_available():
    _GPU_MEM_GB = torch.cuda.get_device_properties(0).total_memory / 1024 ** 3
SKIP_EMPTY_CACHE = _GPU_MEM_GB > 80

def maybe_empty_cache():
    if not SKIP_EMPTY_CACHE and torch.cuda.is_available():
        torch.cuda.empty_cache()