diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py
index 9543ddd..eaab16c 100644
--- a/laborious/utils/repository/model_repository.py
+++ b/laborious/utils/repository/model_repository.py
@@ -705,7 +705,11 @@ class MLFlowRepository(SientiaMonitoring):
     """
 
     def get_prediction_data(
-        self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str, predict_flavor: str
+        self,
+        prediction_model: Any,
+        retrain_dataset: pd.DataFrame,
+        target_name: str,
+        predict_flavor: str,
     ) -> pd.DataFrame:
         """
         Get prediction data from prediction model.
@@ -861,7 +865,9 @@ class MLFlowRepository(SientiaMonitoring):
         prediction_model.fit(retrain_dataset)
 
         # get prediction data
-        prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name, predict_flavor)
+        prediction_data = self.get_prediction_data(
+            prediction_model, retrain_dataset, target_name, predict_flavor
+        )
 
         self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
 
diff --git a/requirements.txt b/requirements.txt
index abc1753..3921162 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,7 +4,7 @@ sqlalchemy
 asyncua
 redis
 git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1
-git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.5
+git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.6
 prometheus-client
 botocore
 boto3
diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py
index a86c7ec..466fae5 100644
--- a/tests/laborious/utils/repository/test_model_repository.py
+++ b/tests/laborious/utils/repository/test_model_repository.py
@@ -865,7 +865,10 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
     prediction_model.fit.assert_called_once_with(pd_merge.return_value)
 
     mlflow_repository.get_prediction_data.assert_called_once_with(
-        prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable, 'pyfunc'
+        prediction_model,
+        pd_merge.return_value,
+        data_model.fit.return_value.target_variable,
+        'pyfunc',
     )
 
     assert output == {