diff --git a/src/gretel_trainer/benchmark/gretel/models.py b/src/gretel_trainer/benchmark/gretel/models.py
index d627c602..324860d0 100644
--- a/src/gretel_trainer/benchmark/gretel/models.py
+++ b/src/gretel_trainer/benchmark/gretel/models.py
@@ -8,7 +8,7 @@
 
 from gretel_client.projects.exceptions import ModelConfigError
 from gretel_client.projects.models import read_model_config
-from gretel_trainer.benchmark.core import BenchmarkException, Dataset, Datatype
+from gretel_trainer.benchmark.core import BenchmarkException, Dataset
 
 GretelModelConfig = Union[str, Path, dict]