diff --git a/cli/quantize_cpu_naive.py b/cli/quantize_cpu_naive.py index ff4992f..727ce22 100644 --- a/cli/quantize_cpu_naive.py +++ b/cli/quantize_cpu_naive.py @@ -5,9 +5,10 @@ import os import psutil import torch.backends.quantized import transformers -from hivemind.utils.logging import get_logger +from hivemind.utils.logging import get_logger, use_hivemind_log_handler from tqdm.auto import trange +use_hivemind_log_handler("in_root_logger") logger = get_logger(__file__) DTYPE_MAP = dict(bfloat16=torch.bfloat16, float16=torch.float16, float32=torch.float32, auto="auto")