Hi,

You can use Databricks Container Services on clusters with GPUs to create portable deep learning environments with customized libraries. See Customize containers with Databricks Container Services for instructions.

To create custom images for GPU clusters, you must select a standard runtime version instead of Databricks Runtime ML for GPU. When you select Use your own Docker container, you can choose GPU clusters with a standard runtime version. The custom images for GPU clusters are based on the official CUDA containers, which is different from Databricks Runtime ML for GPU.

When you create custom images for GPU clusters, you cannot change the NVIDIA driver version, because it must match the driver version on the host machine.

 Docker Hub contains example base images with GPU capability. The Dockerfiles used to generate these images are located in the example containers GitHub repository, which also has details on what the example images provide and how to customize them.

Please refer to : https://docs.databricks.com/clusters/gpu.html#databricks-container-services-on-gpu-clusters

Please let us know if this helps. 

Also please tag @Debayan​ with your next response which will notify me, Thank you!