ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-04-08 06:24:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

CUDA_VISIBLE_DEVICES should not be set to "all"

Browse Source
This commit is contained in:
allegroai 2020-05-31 14:00:51 +03:00
parent 22c5f043aa
commit 5f112447f7
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
trains_agent/session.py
View File

@ -78,7 +78,12 @@ class Session(_Session):
os.environ['CUDA_VISIBLE_DEVICES'] = os.environ['NVIDIA_VISIBLE_DEVICES'] = 'none'
if kwargs.get('gpus') and not os.environ.get('KUBERNETES_SERVICE_HOST') \
and not os.environ.get('KUBERNETES_PORT'):
os.environ['CUDA_VISIBLE_DEVICES'] = os.environ['NVIDIA_VISIBLE_DEVICES'] = kwargs.get('gpus')
# CUDA_VISIBLE_DEVICES does not support 'all'
if kwargs.get('gpus') == 'all':
os.environ.pop('CUDA_VISIBLE_DEVICES', None)
os.environ['NVIDIA_VISIBLE_DEVICES'] = kwargs.get('gpus')
else:
os.environ['CUDA_VISIBLE_DEVICES'] = os.environ['NVIDIA_VISIBLE_DEVICES'] = kwargs.get('gpus')
if kwargs.get('only_load_config'):
from trains_agent.backend_api.config import load
self.config = load()