diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py
index 888b5c8..c3ef12d 100644
--- a/laborious/utils/repository/model_repository.py
+++ b/laborious/utils/repository/model_repository.py
@@ -138,6 +138,7 @@ class MLFlowRepository():
             treated_data, y,
             left_index=True,
             right_index=True,
+            on='id',
             how='left',           # or 'inner', depending on your needs
             validate='one_to_one'  # ensures each index appears only once in both
         )