# Create a device request asking for all GPUs (equivalent to --gpus all)
gpu_request = [DeviceRequest(count=-1, capabilities=[['gpu']])]

self._container = await asyncio.to_thread(
    client.containers.create,
    self._image,
    name=self.container_name,
    entrypoint=shell_command,
    command=command,
    tty=True,
    detach=True,
    auto_remove=self._auto_remove,
    volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes},
    working_dir="/workspace",
    extra_hosts=self._extra_hosts,
    # Pass the GPU request to the container creation call
    device_requests=gpu_request,
)
await asyncio.to_thread(self._container.start)