diff --git a/modules/ml-pipeline/src/pipeline/core/MLModels.py b/modules/ml-pipeline/src/pipeline/core/MLModels.py
index 4cf8b08..4fc572a 100644
--- a/modules/ml-pipeline/src/pipeline/core/MLModels.py
+++ b/modules/ml-pipeline/src/pipeline/core/MLModels.py
@@ -149,6 +149,8 @@ class AutogluonAutoML:
         "time_limit",
         "presets",
         "excluded_model_types",
+        "infer_limit",
+        "infer_limit_batch_size",
     ]
 
     def load_model(self, path: Union[Path, str]) -> None:
@@ -203,6 +205,8 @@ class AutogluonAutoML:
             time_limit=model_hyperparameters["time_limit"],
             presets=model_hyperparameters["presets"],
             excluded_model_types=model_hyperparameters["excluded_model_types"],
+            infer_limit=model_hyperparameters["infer_limit"],
+            infer_limit_batch_size=model_hyperparameters["infer_limit_batch_size"],
         )
 
     def predict(