Unset CUDA_VISIBLE_DEVICES for worker processes (aws#122)

* Unset CUDA_VISIBLE_DEVICES for worker processes

* Add comments

Author: icywang86rui

src/sagemaker_tensorflow_container/training.py

@@ -202,6 +202,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)
 <<<<<<< HEAD
     if install_module: