---
- name: Setup NVIDIA drivers
  become: true
  tasks:
   - name: "Adding NVIDIA CUDA repository"
     ansible.builtin.apt_repository:
      repo: "deb [signed-by=/usr/share/keyrings/cuda-archive-keyring.gpg] https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
      filename: nvidia-upstream-ppa
   - name: "Installing CUDA"
     ansible.builtin.apt:
      pkg:
       - cuda-12
       - nvidia-headless
   - name: "Making sure NVIDIA UVM is working"
     community.general.modprobe:
      name: "nvidia_uvm"
      state: present