diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py
index 395865e..9543ddd 100644
--- a/laborious/utils/repository/model_repository.py
+++ b/laborious/utils/repository/model_repository.py
@@ -705,14 +705,17 @@ class MLFlowRepository(SientiaMonitoring):
     """
 
     def get_prediction_data(
-        self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str
+        self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str, predict_flavor: str
     ) -> pd.DataFrame:
         """
         Get prediction data from prediction model.
         """
         input_index = retrain_dataset.index
 
-        prediction_data = prediction_model.predict(retrain_dataset)
+        if predict_flavor == 'pyfunc':
+            prediction_data = prediction_model.predict({}, retrain_dataset)
+        else:
+            prediction_data = prediction_model.predict(retrain_dataset)
 
         if isinstance(prediction_data, pd.DataFrame):
             prediction_data.columns = pd.Index(['prediction'])
@@ -858,7 +861,7 @@ class MLFlowRepository(SientiaMonitoring):
         prediction_model.fit(retrain_dataset)
 
         # get prediction data
-        prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name)
+        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/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py
index 40782fe..a86c7ec 100644
--- a/tests/laborious/utils/repository/test_model_repository.py
+++ b/tests/laborious/utils/repository/test_model_repository.py
@@ -865,7 +865,7 @@ 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
+        prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable, 'pyfunc'
     )
 
     assert output == {
@@ -951,7 +951,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
     prediction_model.fit.assert_called_once_with(transformed_data)
 
     mlflow_repository.get_prediction_data.assert_called_once_with(
-        prediction_model, transformed_data, 'feat_1'
+        prediction_model, transformed_data, 'feat_1', 'pyfunc'
     )
 
     assert output == {
@@ -1494,8 +1494,11 @@ def test_get_prediction_data_dataframe(mlflow_repository):
     retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
     prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
     target_name = 'target'
+    predict_flavor = 'sklearn'
 
-    result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
 
     prediction_model.predict.assert_called_once_with(retrain_dataset)
     assert 'prediction' in result.columns
@@ -1509,11 +1512,32 @@ def test_get_prediction_data_array(mlflow_repository):
     retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
     prediction_model.predict.return_value = [5, 6]
     target_name = 'target'
+    predict_flavor = 'sklearn'
 
-    result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
 
     prediction_model.predict.assert_called_once_with(retrain_dataset)
     assert 'prediction' in result.columns
     assert 'target' in result.columns
     assert 'timestamp' in result.columns
     assert result.index.tolist() == [0, 1]
+
+
+def test_get_prediction_data_pyfunc(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
+    target_name = 'target'
+    predict_flavor = 'pyfunc'
+
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
+
+    prediction_model.predict.assert_called_once_with({}, retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]