sudo apt install cuda-toolkit-12-6 cuda-compat-12-6

# Edit the file /etc/profile.d/cuda.sh and add on the last line:
export LD_LIBRARY_PATH=/usr/local/cuda/compat${LD_LIBRARY_PATH:+:${LD_LIBRAY_PATH}}

# Save the file, log out of your VM and log back in. That should solve the issue.

ubuntu@gputest:~$ lspci | grep NVIDIA 
00:05.0 3D controller: NVIDIA Corporation Device 20f1 (rev a1)

ubuntu@gputest:~$ journalctl -u nvidia-gridd | tail 
Sep 04 08:54:23 jammy-gpu systemd[1]: Starting NVIDIA Grid Daemon...
Sep 04 08:54:24 jammy-gpu nvidia-gridd[724]: Started (724)
Sep 04 08:54:24 jammy-gpu systemd[1]: Started NVIDIA Grid Daemon.
Sep 04 08:54:24 jammy-gpu nvidia-gridd[724]: Configuration parameter ( ServerAddress  ) not set
Sep 04 08:54:24 jammy-gpu nvidia-gridd[724]: vGPU Software package (0)
Sep 04 08:54:24 jammy-gpu nvidia-gridd[724]: Ignore service provider and node-locked licensing
Sep 04 08:54:24 jammy-gpu nvidia-gridd[724]: NLS initialized
Sep 04 08:54:24 jammy-gpu nvidia-gridd[724]: Acquiring license. (Info: nvidiadls02.it.ntnu.no; NVIDIA Virtual Compute Server)
Sep 04 08:54:27 jammy-gpu nvidia-gridd[724]: License acquired successfully. (Info: nvidiadls02.it.ntnu.no, NVIDIA Virtual Compute Server; Expiry: 2023-9-5 8:54:16 GMT)

ubuntu@gputest:~$ nvidia-smi 
Mon Sep  4 08:58:11 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.125.06   Driver Version: 525.125.06   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GRID V100D-8C       On   | 00000000:00:05.0 Off |                  N/A |
| N/A   N/A    P0    N/A /  N/A |      0MiB /  8192MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

$ git clone https://github.com/NVIDIA/cuda-samples.git

ubuntu@jammy-gpu:~$ cd cuda-samples/Samples/1_Utilities/deviceQuery
ubuntu@jammy-gpu:~/cuda-samples/Samples/1_Utilities/deviceQuery$ make
  ... lots-of-text-from-make ... 
ubuntu@jammy-gpu:~/cuda-samples/Samples/1_Utilities/deviceQuery$ ./deviceQuery 
./deviceQuery Starting...

 CUDA Device Query (Runtime API) version (CUDART static linking)

Detected 1 CUDA Capable device(s)

Device 0: "GRID V100D-8C"
  CUDA Driver Version / Runtime Version          12.0 / 12.0
  CUDA Capability Major/Minor version number:    7.0
  Total amount of global memory:                 8192 MBytes (8589934592 bytes)
  (080) Multiprocessors, (064) CUDA Cores/MP:    5120 CUDA Cores
  GPU Max Clock rate:                            1380 MHz (1.38 GHz)
  Memory Clock rate:                             877 Mhz
  Memory Bus Width:                              4096-bit
  L2 Cache Size:                                 6291456 bytes
  Maximum Texture Dimension Size (x,y,z)         1D=(131072), 2D=(131072, 65536), 3D=(16384, 16384, 16384)
  Maximum Layered 1D Texture Size, (num) layers  1D=(32768), 2048 layers
  Maximum Layered 2D Texture Size, (num) layers  2D=(32768, 32768), 2048 layers
  Total amount of constant memory:               65536 bytes
  Total amount of shared memory per block:       49152 bytes
  Total shared memory per multiprocessor:        98304 bytes
  Total number of registers available per block: 65536
  Warp size:                                     32
  Maximum number of threads per multiprocessor:  2048
  Maximum number of threads per block:           1024
  Max dimension size of a thread block (x,y,z): (1024, 1024, 64)
  Max dimension size of a grid size    (x,y,z): (2147483647, 65535, 65535)
  Maximum memory pitch:                          2147483647 bytes
  Texture alignment:                             512 bytes
  Concurrent copy and kernel execution:          Yes with 7 copy engine(s)
  Run time limit on kernels:                     No
  Integrated GPU sharing Host Memory:            No
  Support host page-locked memory mapping:       Yes
  Alignment requirement for Surfaces:            Yes
  Device has ECC support:                        Disabled
  Device supports Unified Addressing (UVA):      Yes
  Device supports Managed Memory:                No
  Device supports Compute Preemption:            Yes
  Supports Cooperative Kernel Launch:            Yes
  Supports MultiDevice Co-op Kernel Launch:      Yes
  Device PCI Domain ID / Bus ID / location ID:   0 / 0 / 5
  Compute Mode:
     < Default (multiple host threads can use ::cudaSetDevice() with device simultaneously) >

deviceQuery, CUDA Driver = CUDART, CUDA Driver Version = 12.0, CUDA Runtime Version = 12.0, NumDevs = 1
Result = PASS

curl https://get.docker.com | sh && sudo systemctl --now enable docker

# Enable the repositories
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# Install the package
sudo apt update && sudo apt -y install nvidia-docker2

# Restart the docker daemon
sudo systemctl restart docker

# Run a test to verifiy that it works
sudo docker run --rm --gpus all nvidia/cuda:12.9.1-cudnn-runtime-ubuntu24.04 nvidia-smi

# Optionally run a test with Tensorflow that actually runs a bit of code on the GPU via docker
sudo docker run --gpus all -it --rm tensorflow/tensorflow:2.14.0-gpu \
   python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"