fix wandb so mypy doesn't complain

axolotl-ai-cloud · Sep 13, 2023 · 77437f0 · 77437f0
1 parent a1d4fdc
commit 77437f0
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/scripts/finetune.py b/scripts/finetune.py
@@ -26,7 +26,7 @@
 from axolotl.utils.distributed import is_main_process
 from axolotl.utils.models import load_tokenizer
 from axolotl.utils.tokenization import check_dataset_labels
-from axolotl.utils.wandb import setup_wandb_env_vars
+from axolotl.utils.wandb_ import setup_wandb_env_vars
 
 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 src_dir = os.path.join(project_root, "src")

diff --git a/src/axolotl/utils/callbacks.py b/src/axolotl/utils/callbacks.py
@@ -11,6 +11,7 @@
 import pandas as pd
 import torch
 import torch.distributed as dist
+import wandb
 from datasets import load_dataset
 from optimum.bettertransformer import BetterTransformer
 from tqdm import tqdm
@@ -24,7 +25,6 @@
 )
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy
 
-import wandb
 from axolotl.utils.bench import log_gpu_memory_usage
 from axolotl.utils.distributed import (
     barrier,

diff --git a/src/axolotl/utils/wandb.py → src/axolotl/utils/wandb_.py b/src/axolotl/utils/wandb.py → src/axolotl/utils/wandb_.py