diff --git a/smtb/train.py b/smtb/train.py
index 1dc7c97..40aa885 100644
--- a/smtb/train.py
+++ b/smtb/train.py
@@ -1,11 +1,11 @@
 import argparse
 from pathlib import Path
 
+import wandb
 from lightning.pytorch import Trainer, seed_everything
 from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, RichProgressBar
 from lightning.pytorch.loggers import WandbLogger
 
-import wandb
 from smtb.data import DownstreamDataModule
 from smtb.model import RegressionModel