diff --git a/pyproject.toml b/pyproject.toml
index 243ecd87..470842b5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -90,6 +90,7 @@ ignore = [
 pydocstyle.convention = "google"
 isort.required-imports = ["from __future__ import annotations"]
 isort.split-on-trailing-comma = false
+isort.known-third-party = ["wandb"]
 
 [tool.ruff.format]
 docstring-code-format = true
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index eb0e8a86..9403534c 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -5,9 +5,9 @@
 import numpy as np
 import pytest
 import torch
+import wandb
 from pymatgen.core import Lattice, Structure
 
-import wandb
 from chgnet.data.dataset import StructureData, get_train_val_test_loader
 from chgnet.model import CHGNet
 from chgnet.trainer import Trainer