diff --git a/deepof/train_utils.py b/deepof/train_utils.py
index c75e8c56786ee4115cf3032262d131e8041e9912..f90c62d8398315b3aa04d1f884388fb21c61f6af 100644
--- a/deepof/train_utils.py
+++ b/deepof/train_utils.py
@@ -234,11 +234,4 @@ def tune_search(
 
     print(tuner.results_summary())
 
-    return best_hparams, best_run
-
-
-# TODO:
-#    - load_treatments should be part of the main data module. If available in the main directory,
-#    a table (preferrable in csv) should be loaded as metadata of the coordinates automatically.
-#    This becomes particularly important por the supervised models that include phenotype classification
-#    alongside the encoding.
+    return best_hparams, best_run
\ No newline at end of file