Use Nvidia GPU in AWS Batch

For a job, I run puppeteer in the headful mode in AWS Batch. I found a docker file that uses the image node:16.5-slim, then downloads Chrome, and then uses Xvfb.

If I run it on an EC2 instance with GPU, Chrome still uses the integrated video card (llvmpipe) not taking advantage of the powerful Nvidia GPU.

I found online a couple of attempts, one I heard they had to switch to Windows (not supported by AWS Batch). They also used an AMI only available on EC2: https://stackoverflow.com/questions/67813309/run-headful-chromium-with-nvidia-gpu-and-xvfb-on-ec2-instance-with-prime

And the other solution also doesn't seem to be suitable because of its complexity to set up: https://boygiandi.medium.com/ch%E1%BA%A1y-chrome-v%E1%BB%9Bi-gpu-tr%C3%AAn-instance-ec2-aws-69369d73235a

I wonder if there is a docker image I could use as a starting point where the Nvidia GPU is already set up and Chrome could take advantage of it with only a few extra configurations. Thanks!