From 77c43911a2c118a865b0bd28313b18d1f3d3a61c Mon Sep 17 00:00:00 2001 From: Marek Novotny Date: Mon, 13 Nov 2023 16:14:02 +0100 Subject: [PATCH] Add flavor as self reference --- h2o-py-mlflow-flavor/h2o_mlflow_flavor/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/h2o-py-mlflow-flavor/h2o_mlflow_flavor/__init__.py b/h2o-py-mlflow-flavor/h2o_mlflow_flavor/__init__.py index cdbfc0fbc511..4d6aa93bf459 100644 --- a/h2o-py-mlflow-flavor/h2o_mlflow_flavor/__init__.py +++ b/h2o-py-mlflow-flavor/h2o_mlflow_flavor/__init__.py @@ -35,6 +35,7 @@ _validate_and_prepare_target_save_path, ) from mlflow.tracking.artifact_utils import _download_artifact_from_uri +from mlflow.utils.requirements_utils import _get_pinned_requirement _logger = logging.getLogger(__name__) @@ -47,7 +48,7 @@ def get_default_pip_requirements(): Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment that, at minimum, contains these requirements. """ - return [] + return [_get_pinned_requirement("h2o_mlflow_flavor")] def get_default_conda_env():