Commit 6f034645 authored by Daniele Venzano's avatar Daniele Venzano

Set env variable needed for tensorflow gpu

parent 81fb19cb
......@@ -83,6 +83,7 @@ def goog_tensorflow_service(memory_limit, core_limit, image, gpu=False):
if gpu:
service['labels'] = ['gpu']
service['environment'] = [['NVIDIA_VISIBLE_DEVICES', 'all']]
return service
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment