diff --git a/rubicon_ml/client/mixin.py b/rubicon_ml/client/mixin.py index f71efc96..fb362162 100644 --- a/rubicon_ml/client/mixin.py +++ b/rubicon_ml/client/mixin.py @@ -274,13 +274,7 @@ def log_h2o_model( with tempfile.TemporaryDirectory() as temp_dir_name: if use_mojo: model_data_path = f"{temp_dir_name}/{artifact_name}.zip" - h2o_model.save_mojo(path=model_data_path) - - artifact = self.log_artifact( - name=artifact_name, - data_directory=model_data_path, - **log_artifact_kwargs, - ) + h2o_model.download_mojo(path=model_data_path) else: model_data_path = h2o.save_model( h2o_model, @@ -289,11 +283,11 @@ def log_h2o_model( path=temp_dir_name, ) - artifact = self.log_artifact( - name=artifact_name, - data_path=model_data_path, - **log_artifact_kwargs, - ) + artifact = self.log_artifact( + name=artifact_name, + data_path=model_data_path, + **log_artifact_kwargs, + ) return artifact