diff --git a/openqdc/utils/regressor.py b/openqdc/utils/regressor.py
index f9d5199..a980ef5 100644
--- a/openqdc/utils/regressor.py
+++ b/openqdc/utils/regressor.py
@@ -7,8 +7,6 @@
 import pandas as pd
 from loguru import logger
 
-from openqdc.datasets.base import BaseDataset
-
 
 def non_nan_idxs(array):
     """
@@ -119,7 +117,7 @@ def __init__(
         self._post_init()
 
     @classmethod
-    def from_openqdc_dataset(cls, dataset: BaseDataset, *args, **kwargs) -> "Regressor":
+    def from_openqdc_dataset(cls, dataset, *args, **kwargs) -> "Regressor":
         """
         Initialize the regressor object from an openqdc dataset. This is the default method.
         *args and and **kwargs are passed to the __init__ method and depends on the specific regressor.