Commit 5f37a5a2 authored by Tom Ward's avatar Tom Ward Committed by Copybara-Service
Browse files

No longer use `runtime=nvidia` to expose GPUs to the container.

Instead use `DeviceRequest`, which is equivalent to running with `--gpu all`.

PiperOrigin-RevId: 431933575
Change-Id: Iee9f70ae8f291d16850e4f58349559d9468ef4d8
parent 552586eb
......@@ -227,10 +227,14 @@ def main(argv):
])
client = docker.from_env()
device_requests = [
docker.types.DeviceRequest(driver='nvidia', capabilities=[['gpu']])
] if FLAGS.use_gpu else None
container = client.containers.run(
image=FLAGS.docker_image_name,
command=command_args,
runtime='nvidia' if FLAGS.use_gpu else None,
device_requests=device_requests,
remove=True,
detach=True,
mounts=mounts,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment