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()