Specify driver and cuda version when using saas-linux-medium-amd64-gpu-standard

Is it possible to specify cuda and nvidia driver version when using saas-linux-medium-amd64-gpu-standard

The drivers being provided by host are too old and I am getting:

tensorflow.python.framework.errors_impl.InternalError: cudaGetDevice() failed. Status: CUDA driver version is insufficient for CUDA runtime version

image: nvidia/cuda:12.3.0-base-ubuntu22.04

This is still loading up cuda 11.4 driver.

@gabrielengel_gl Maybe you can help?