Save also mojo

h2oai · Nov 15, 2024 · de91bca · de91bca
1 parent 14021df
commit de91bca
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/h2o-sagemaker/gbm/gbm_scripts/train b/h2o-sagemaker/gbm/gbm_scripts/train
@@ -126,6 +126,7 @@ def _train_model(hyperparameters={}):
     training_params = json.loads(training_params_str)
     output_path = os.path.join(prefix, 'output')
     model_path = os.path.join(prefix, 'model')
+    mojo_path = model_path
 
     gbm_params['distribution'] = training_params['distribution']
 
@@ -175,6 +176,8 @@ def _train_model(hyperparameters={}):
 
         model_path = h2o.save_model(gbm_model, path=model_path)
         print("Model Path with Name: {}".format(model_path))
+        mojo_path = gbm_model.save_mojo(path=mojo_path)
+        print("Mojo Path with Name: {}".format(mojo_path))
         print('Training complete.')
     except Exception as e:
         # Write out an error file. This will be returned as the failureReason