diff --git a/cli/train.py b/cli/train.py
index bc5e287..9ed1cd0 100644
--- a/cli/train.py
+++ b/cli/train.py
@@ -4,6 +4,7 @@ from omegaconf import DictConfig
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
 from pytorch_lightning.loggers import MLFlowLogger
 
+from enhancer.data.dataset import EnhancerDataset
 
 @hydra.main(config_path="train_config",config_name="config")
 def main(config: DictConfig):