From 0f96bc55f13135279e682fbd4fe4bf821ccb2233 Mon Sep 17 00:00:00 2001 From: Michael Duong Date: Sun, 22 Oct 2023 21:05:07 +0000 Subject: [PATCH] add time to inference to model --- modules/ml-pipeline/src/pipeline/core/MLModels.py | 4 ++++ 1 file changed, 4 insertions(+) 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(