diff --git a/modules/ml-pipeline/src/pipeline/configs/build_model.yaml b/modules/ml-pipeline/src/pipeline/configs/build_model.yaml
index d296e6a..1ebb62d 100644
--- a/modules/ml-pipeline/src/pipeline/configs/build_model.yaml
+++ b/modules/ml-pipeline/src/pipeline/configs/build_model.yaml
@@ -16,3 +16,5 @@ default:
       time_limit: 4000
       presets: medium_quality
       excluded_model_types: ['KNN', 'RF']
+      infer_limit: 0.05
+      infer_limit_batch_size: 10000
diff --git a/modules/ml-pipeline/src/pipeline/core/Logger.py b/modules/ml-pipeline/src/pipeline/core/Logger.py
index 2194063..d2f6c61 100644
--- a/modules/ml-pipeline/src/pipeline/core/Logger.py
+++ b/modules/ml-pipeline/src/pipeline/core/Logger.py
@@ -21,7 +21,6 @@ def setup_logger():
 
     # Add the stream handler to the logger
     logger.addHandler(stream_handler)
-
     logger.propagate = False
 
     return logger
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(
diff --git a/modules/ml-pipeline/src/pipeline/dvc.lock b/modules/ml-pipeline/src/pipeline/dvc.lock
index 16eb857..20dd532 100644
--- a/modules/ml-pipeline/src/pipeline/dvc.lock
+++ b/modules/ml-pipeline/src/pipeline/dvc.lock
@@ -5,8 +5,8 @@ stages:
     deps:
     - path: 1_prepare_data.py
       hash: md5
-      md5: c9f030df733e318b80d1fa91b7732f79
-      size: 5132
+      md5: 896d3d88a4a9f68d174efe71dc089517
+      size: 4222
     params:
       configs/settings.yaml:
         default.feature_processor.feature_processor_config.drop_columns:
@@ -29,20 +29,20 @@ stages:
     outs:
     - path: data/prepared_data/
       hash: md5
-      md5: cd75be9fecff0c647792dd2db648085c.dir
-      size: 37056053
+      md5: 6bfdb621b608648c017bf2323f7b5052.dir
+      size: 37048968
       nfiles: 2
   build_model:
     cmd: python 2_build_model.py
     deps:
     - path: 2_build_model.py
       hash: md5
-      md5: 84699d208874c52accaff61c6af9bb0a
-      size: 5359
+      md5: 7b79f280b8b0d5bc6f07669e7cc37c6a
+      size: 4150
     - path: data/prepared_data
       hash: md5
-      md5: cd75be9fecff0c647792dd2db648085c.dir
-      size: 37056053
+      md5: 6bfdb621b608648c017bf2323f7b5052.dir
+      size: 37048968
       nfiles: 2
     params:
       configs/build_model.yaml:
@@ -63,32 +63,34 @@ stages:
               excluded_model_types:
               - KNN
               - RF
+              infer_limit: 0.05
+              infer_limit_batch_size: 10000
     outs:
     - path: data/model/
       hash: md5
-      md5: 7a5527f779efcb1a7db068148b6bcc45.dir
-      size: 422448184
+      md5: f2999107de7572ea5ff0f2d774fa83b8.dir
+      size: 424943352
       nfiles: 27
     - path: metrics/fit_metrics.json
       hash: md5
-      md5: 77790bb9485c04c77125e361921c3774
-      size: 225
+      md5: 9537e7ebc2eb32b421a7cabd2005f00b
+      size: 223
   generate_predictions:
     cmd: python 3_generate_predictions.py
     deps:
     - path: 3_generate_predictions.py
       hash: md5
-      md5: 5ef2856a5a977304f1ec01f9b4205262
-      size: 3028
+      md5: 0a70ad4dfe99414a75d1261c75a177b9
+      size: 2464
     - path: data/model
       hash: md5
-      md5: 7a5527f779efcb1a7db068148b6bcc45.dir
-      size: 422448184
+      md5: f2999107de7572ea5ff0f2d774fa83b8.dir
+      size: 424943352
       nfiles: 27
     - path: data/prepared_data
       hash: md5
-      md5: cd75be9fecff0c647792dd2db648085c.dir
-      size: 37056053
+      md5: 6bfdb621b608648c017bf2323f7b5052.dir
+      size: 37048968
       nfiles: 2
     params:
       configs/settings.yaml:
@@ -100,25 +102,25 @@ stages:
     outs:
     - path: data/predictions/
       hash: md5
-      md5: 28d2876e6c6d5cc64844ecc1d6ac40b2.dir
-      size: 346687
+      md5: f4439a56669f84bc51a9fcb4cd08353f.dir
+      size: 346539
       nfiles: 1
   generate_metrics:
     cmd: python 4_generate_metrics.py
     deps:
     - path: 4_generate_metrics.py
       hash: md5
-      md5: 2c9fb78955a8c19cff0a098976f81d1b
-      size: 4487
+      md5: 4fedb86d89d528f0a6597934ba3890a0
+      size: 3484
     - path: data/predictions
       hash: md5
-      md5: 28d2876e6c6d5cc64844ecc1d6ac40b2.dir
-      size: 346687
+      md5: f4439a56669f84bc51a9fcb4cd08353f.dir
+      size: 346539
       nfiles: 1
     - path: data/prepared_data
       hash: md5
-      md5: cd75be9fecff0c647792dd2db648085c.dir
-      size: 37056053
+      md5: 6bfdb621b608648c017bf2323f7b5052.dir
+      size: 37048968
       nfiles: 2
     params:
       configs/settings.yaml:
@@ -128,15 +130,15 @@ stages:
     outs:
     - path: metrics/metrics.json
       hash: md5
-      md5: 7afd04d656dc83ad6aa942d9c63f5b4e
+      md5: 357904cf106279be5a578e8faefa5d80
       size: 224
   startup_cleanup:
     cmd: python 0_startup_cleanup.py
     deps:
     - path: 0_startup_cleanup.py
       hash: md5
-      md5: fbb7e3b1b98b517c870f3e1df3e7f695
-      size: 1676
+      md5: b1b12f6b6393fbf8b83d23684df0a3d4
+      size: 1220
     params:
       configs/settings.yaml:
         default.startup_cleanup.artefacts: ./data