fix: logic error in MWMS

Lokiiiiii · Lokiiiiii · commit 4e1f0af4a4bc · 2022-06-03T19:27:12.000-07:00
diff --git a/src/sagemaker_tensorflow_container/training.py b/src/sagemaker_tensorflow_container/training.py
@@ -183,8 +183,8 @@ def train(env, cmd_args):
 
     elif multi_worker_mirrored_strategy_enabled:
 
-        env_vars["TF_CONFIG"] = _build_tf_config_for_mwms(
-            hosts=env.hosts, current_host=env.current_host
+        env_vars["TF_CONFIG"] = json.dumps(
+            _build_tf_config_for_mwms(hosts=env.hosts, current_host=env.current_host)
         )
         logger.info("Running distributed training job with multi_worker_mirrored_strategy setup")
 

Original file line number	Diff line number	Diff line change
`@@ -183,8 +183,8 @@ def train(env, cmd_args):`
`183`	`183`
`184`	`184`	`elif multi_worker_mirrored_strategy_enabled:`
`185`	`185`
`186`		`- env_vars["TF_CONFIG"] = _build_tf_config_for_mwms(`
`187`		`- hosts=env.hosts, current_host=env.current_host`
	`186`	`+ env_vars["TF_CONFIG"] = json.dumps(`
	`187`	`+ _build_tf_config_for_mwms(hosts=env.hosts, current_host=env.current_host)`
`188`	`188`	`)`
`189`	`189`	`logger.info("Running distributed training job with multi_worker_mirrored_strategy setup")`
`190`	`190`