fix: download_mojo instead of save_mojo

capitalone · Sep 26, 2024 · dd54552 · dd54552
1 parent c95139e
commit dd54552
Showing 1 changed file with 6 additions and 12 deletions.
diff --git 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