NVIDIA cuda version for docker

Hey folks, need some help here! We’re trying to run an AWS batch job that uses an Nvidia GPU in darts model definition. It works fine on our local machine, but when we containerize it with Docker, it fails to detect the GPU available on the instance.
Host-machine has cuda 11.4 nvidia drivers pre- installed and nvidia toolkit of the same version. How can I build a docker image to use these nvidia drivers so that it detects the GPU in the instance?