diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py
index 4e8e6c3..d01b53b 100644
--- a/chebai_graph/preprocessing/datasets/base.py
+++ b/chebai_graph/preprocessing/datasets/base.py
@@ -188,7 +188,7 @@ def _after_setup(self, **kwargs) -> None:
         self._setup_properties()
         super()._after_setup(**kwargs)
 
-    def _preprocess_smiles_for_pred(
+    def _preprocess_molecule_for_pred(
         self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None
     ) -> Optional[dict]:
         """Preprocess prediction data."""