diff --git a/ansible/roles/installation/nvidia-docker/tasks/main.yml b/ansible/roles/installation/nvidia-docker/tasks/main.yml
index ce9acd790..70352ca67 100644
--- a/ansible/roles/installation/nvidia-docker/tasks/main.yml
+++ b/ansible/roles/installation/nvidia-docker/tasks/main.yml
@@ -37,15 +37,22 @@
       state: present
     become: yes
 
-  - name: query prime-select
-    command: prime-select query
-    become: yes
-    register: prime_select_results
+  - name: Query and set prime-select
+    block:
+    - name: query prime-select
+      command: prime-select query
+      become: yes
+      register: prime_select_results
 
-  - name: switch driver to nvidia
-    command: prime-select nvidia
-    become: yes
-    when: "prime_select_results.stdout != 'nvidia'"
+    - name: switch driver to nvidia
+      command: prime-select nvidia
+      become: yes
+      when: "prime_select_results.stdout != 'nvidia'"
+
+    - name: query prime-select again, check it's been set
+      command: prime-select query
+      become: yes
+      register: prime_select_results
 
   - name: end playbook so user can reboot
     fail:
diff --git a/ansible/roles/installation/nvidia-docker/tasks/nvidia-container-toolkit.yml b/ansible/roles/installation/nvidia-docker/tasks/nvidia-container-toolkit.yml
index 8418887bc..c74fb4f4d 100644
--- a/ansible/roles/installation/nvidia-docker/tasks/nvidia-container-toolkit.yml
+++ b/ansible/roles/installation/nvidia-docker/tasks/nvidia-container-toolkit.yml
@@ -18,46 +18,28 @@
     name: nvidia-modprobe
     state: present
 
-- name: Install nvidia-docker gpg key
-  apt_key:
-    url: https://nvidia.github.io/nvidia-docker/gpgkey
-    state: present
-
-- name: Set ubuntu_version to 18.04 for nvidia-container-toolkit install
-  set_fact:
-    ubuntu_version: "18.04"
-  when:
-    - ansible_distribution|string == 'Ubuntu'
-    - ansible_distribution_release|string == 'bionic'
-
-- name: Set ubuntu_version to 20.04 for nvidia-container-toolkit install
-  set_fact:
-    ubuntu_version: "20.04"
-  when:
-    - ansible_distribution|string == 'Ubuntu'
-    - ansible_distribution_release|string == 'focal'
-
-- name: Set ubuntu_version to 22.04 for nvidia-container-toolkit install
-  set_fact:
-    ubuntu_version: "22.04"
-  when:
-    - ansible_distribution|string == 'Ubuntu'
-    - ansible_distribution_release|string == 'jammy'
-    
-- name: Add nvidia-docker repository
-  apt_repository:
-    repo: deb https://nvidia.github.io/{{ item }}/ubuntu{{ ubuntu_version }}/amd64 /
-    state: present
-  with_items:
-    - nvidia-docker
-
-- name: Add libnvidia-container repository
-  apt_repository:
-    repo: deb https://nvidia.github.io/{{ item }}/stable/ubuntu{{ ubuntu_version }}/amd64 /
-    state: present
-  with_items:
-    - libnvidia-container
-    - nvidia-container-runtime
+- name: Check nvidia-container-toolkit repository key exists
+  stat: 
+    path: /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
+  register: nvidia_container_toolkit_repo_key_exists
+
+# Inspired by: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html#setting-up-nvidia-container-toolkit
+- name: Get and de-armour gpg key, add repository
+  block:
+    - name: Get armoured gpg key
+      ansible.builtin.get_url:
+        url: https://nvidia.github.io/libnvidia-container/gpgkey
+        dest: /tmp/nvidia-container-toolkit-keyring.asc
+
+    - name: Dearmour gpg key
+      shell: cat /tmp/nvidia-container-toolkit-keyring.asc | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
+      become: true        
+
+    - name: Add nvidia-container-toolkit repository, pointing to the dearmored gpg key (backwards compatable, 18.04 should work with up to at least 22.04)
+      ansible.builtin.apt_repository:
+        repo: "deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://nvidia.github.io/libnvidia-container/stable/ubuntu18.04/$(ARCH) /"
+        state: present
+  when: nvidia_container_toolkit_repo_key_exists.stat.exists == false
 
 - name: Install nvidia-container-toolkit
   apt: