diff --git a/chebai/cli.py b/chebai/cli.py index 41742596..96262447 100644 --- a/chebai/cli.py +++ b/chebai/cli.py @@ -2,7 +2,7 @@ from lightning.pytorch.cli import LightningArgumentParser, LightningCLI -from chebai.preprocessing.datasets import XYBaseDataModule +from chebai.preprocessing.datasets.base import XYBaseDataModule from chebai.trainer.CustomTrainer import CustomTrainer