Unset CUDA_VISIBLE_DEVICES for worker processes (#122)

* Unset CUDA_VISIBLE_DEVICES for worker processes

* Add comments

Author: icywang86rui

src/sagemaker_tensorflow_container/training.py

@@ -107,6 +107,10 @@ def _run_ps(env):
 
 
 def _run_worker(env):
+    # when _run_ps is called CUDA_VISIBLE_DEVICES is set with os.environ. We need to unset it so the worker
+    # process can use the GPUs.
+    if os.environ.get('CUDA_VISIBLE_DEVICES'):
+        del os.environ['CUDA_VISIBLE_DEVICES']
     env_vars = _env_vars_with_tf_config(env, ps_task=False)
     framework.entry_point.run(env.module_dir, env.user_entry_point, env.to_cmd_args(), env_vars)