diff --git a/ms2pip/_utils/ion_mobility.py b/ms2pip/_utils/ion_mobility.py
index 8a3764b..2ee91fa 100644
--- a/ms2pip/_utils/ion_mobility.py
+++ b/ms2pip/_utils/ion_mobility.py
@@ -27,6 +27,6 @@ def add_im_predictions(self, psm_list: PSMList) -> None:
         """Add ion mobility predictions to the PSMList."""
         logger.info("Predicting ion mobility...")
         predictions: pd.Series = self.predict_fn(
-            psm_list, write_output=False, n_jobs=self.processes
+            psm_list, write_output=False, n_jobs=self.processes, ion_mobility=True
         )
         psm_list["ion_mobility"] = predictions.values
diff --git a/ms2pip/_utils/retention_time.py b/ms2pip/_utils/retention_time.py
index 043f24a..a1015a2 100644
--- a/ms2pip/_utils/retention_time.py
+++ b/ms2pip/_utils/retention_time.py
@@ -104,7 +104,7 @@ def _prepare_deeplc_peptide_df(self):
         """
         column_map = {"peptide": "seq", "modifications": "modifications"}
         peprec = peptide_record.to_dataframe(self.psm_list)[column_map.keys()]
-        self.deeplc_pep_df = peprec.rename(columns=column_map)
+        self.deeplc_pep_df = peprec.rename(columns=column_map) #TODO: allow PSMList input
 
     def _run_deeplc(self):
         """Run DeepLC."""
@@ -133,7 +133,7 @@ def add_rt_predictions(self, psm_list):
         """
         self.psm_list = psm_list
 
-        if self.predictor == "deeplc":
+        if self.predictor == "deeplc": #TODO: throw out
             self._predict_deeplc()
         else:
             raise NotImplementedError(self.predictor)
diff --git a/ms2pip/core.py b/ms2pip/core.py
index f3dd5ec..d8999a1 100644
--- a/ms2pip/core.py
+++ b/ms2pip/core.py
@@ -199,7 +199,7 @@ def predict_library(
         logging.disable(logging.CRITICAL)
         yield predict_batch(
             search_space.filter_psms_by_mz(PSMList(psm_list=list(batch))),
-            add_retention_time=add_retention_time,
+            add_retention_time=add_retention_time, #TODO: Perhaps get calibrated DeepLC instance out of predict_batch 
             add_ion_mobility=add_ion_mobility,
             model=model,
             model_dir=model_dir,