switch to logger

This commit is contained in:
Brandon Thomas
2023-05-28 19:09:37 -07:00
parent a36f1c374e
commit a1a73ea6a5
@@ -5,24 +5,24 @@ from logging import getLogger
LOG = getLogger(__name__) LOG = getLogger(__name__)
def clear_cuda_cache(): def clear_cuda_cache():
print('torch.cuda.empty_cache() ...') LOG.debug('torch.cuda.empty_cache() ...')
torch.cuda.empty_cache() torch.cuda.empty_cache()
print('torch.cuda.empty_cache() ... DONE') LOG.debug('torch.cuda.empty_cache() ... DONE')
def turn_off_jit_profiling(): def turn_off_jit_profiling():
print('torch._C._jit_set_profiling_mode(False) ...') LOG.debug('torch._C._jit_set_profiling_mode(False) ...')
torch._C._jit_set_profiling_mode(False) torch._C._jit_set_profiling_mode(False)
print('torch._C._jit_set_profiling_mode(False) ... DONE') LOG.debug('torch._C._jit_set_profiling_mode(False) ... DONE')
def maybe_hacky_fix_before(): def maybe_hacky_fix_before():
LOG.info("maybe_hacky_fix_before()") LOG.debug("maybe_hacky_fix_before()")
if os.environ.get('FIX_CACHE_BEFORE'): if os.environ.get('FIX_CACHE_BEFORE'):
clear_cuda_cache() clear_cuda_cache()
if os.environ.get('FIX_JIT_BEFORE'): if os.environ.get('FIX_JIT_BEFORE'):
turn_off_jit_profiling() turn_off_jit_profiling()
def maybe_hacky_fix_after(): def maybe_hacky_fix_after():
LOG.info("maybe_hacky_fix_after()") LOG.debug("maybe_hacky_fix_after()")
if os.environ.get('FIX_CACHE_AFTER'): if os.environ.get('FIX_CACHE_AFTER'):
clear_cuda_cache() clear_cuda_cache()
if os.environ.get('FIX_JIT_AFTER'): if os.environ.get('FIX_JIT_AFTER'):