Update

awslabs · Nov 22, 2023 · 358f5bb · 358f5bb
1 parent 7729a8e
commit 358f5bb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/python/graphstorm/sagemaker/sagemaker_train.py b/python/graphstorm/sagemaker/sagemaker_train.py
@@ -228,7 +228,7 @@ def run_train(args, unknownargs):
     yaml_path = download_yaml_config(train_yaml_s3,
         data_path, sagemaker_session)
     graph_config_path = download_graph(graph_data_s3, graph_name,
-        host_rank,  world_size, data_path, sagemaker_session)
+        host_rank, world_size, data_path, sagemaker_session)
     if model_checkpoint_s3 is not None:
         # Download Saved model checkpoint to resume
         download_model(model_checkpoint_s3, restore_model_path, sagemaker_session)

diff --git a/sagemaker/launch/launch_train.py b/sagemaker/launch/launch_train.py
@@ -62,13 +62,11 @@ def run_job(input_args, image, unknowargs):
 
     prefix = f"gs-train-{graph_name}"
 
-    params = {
-        "graph-data-s3": graph_data_s3,
-        "graph-name": graph_name,
-        "model-artifact-s3": model_artifact_s3,
-        "task-type": task_type,
-        "train-yaml-s3": train_yaml_s3,
-    }
+    params = {"task-type": task_type,
+              "graph-name": graph_name,
+              "graph-data-s3": graph_data_s3,
+              "train-yaml-s3": train_yaml_s3,
+              "model-artifact-s3": model_artifact_s3}
     if custom_script is not None:
         params["custom-script"] = custom_script
     if model_checkpoint_to_load is not None: