From dd54552c2c41389403f9f37bff7fa4263723cd1c Mon Sep 17 00:00:00 2001
From: thebrianbn <thebrianbn@gmail.com>
Date: Thu, 26 Sep 2024 14:02:22 -0400
Subject: [PATCH] fix: download_mojo instead of save_mojo

---
 rubicon_ml/client/mixin.py | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

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