GPU support for drone-runner-docker

I am looking for the ability to run builds in docker containers with access to GPU. When starting the runner with --gpus all flag, the docker image created by the runner does not appear to have access to a GPU. Does anyone know how to enable this? New to Go, so digging through the source code is taking me a while. Thanks!

Edit: Runtime seems to be HostConfig parameter, just not clear to me how to set it. https://godoc.org/docker.io/go-docker/api/types/container#HostConfig

If you’re using the NVIDIA Container Runtime, then you should be able to give GPU access to a container by setting the NVIDIA_VISIBLE_DEVICES environment variable.

For example:

---
kind: pipeline
type: docker
name: gpu-tests

platform:
  os: linux
  arch: amd64

node:
  os: linux
  class: gpu

trigger:
  event:
  - push

volumes:
- name: docker-socket
  host:
    path: /var/run/docker.sock

steps:
- name: build-image
  image: docker
  commands:
  - docker build -t my-gpu-image:${DRONE_COMMIT_SHA} .
  volumes:
  - name: docker-socket
    path: /var/run/docker.sock

- name: tests
  image: my-gpu-image:${DRONE_COMMIT_SHA}
  commands:
  - python -c 'import torch; torch.cuda.is_available()'
  environment:
    NVIDIA_VISIBLE_DEVICES: all
  volumes:
  - name: docker-socket
    path: /var/run/docker.sock