diff --git a/Dockerfile b/Dockerfile
index 317b9b09..7fca4370 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,9 +1,15 @@
-FROM nvcr.io/nvidia/tensorrt:24.02-py3
+FROM nvidia/cuda:12.3.2-devel-ubuntu22.04
 
 WORKDIR /rembg
 
 COPY . .
 
+# RUN apt-get update
+# RUN apt install software-properties-common -y
+# RUN add-apt-repository ppa:deadsnakes/ppa -y
+RUN apt-get update
+RUN apt-get install -y python-is-python3 python3 python3-pip
+RUN pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
 RUN python -m pip install ".[gpu,cli]"
 RUN python -c 'from rembg.bg import download_models; download_models()'
 
diff --git a/setup.py b/setup.py
index 3332664e..969e410f 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@
         "twine",
         "wheel",
     ],
-    "gpu": [""],
+    "gpu": ["onnxruntime-gpu"],
     "cli": [
         "aiohttp",
         "asyncer",
@@ -85,7 +85,7 @@
     python_requires=">=3.8, <3.13",
     packages=find_packages(),
     install_requires=install_requires,
-    # dependency_links=['https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/'],
+    dependency_links=['https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/'],
     entry_points=entry_points,
     extras_require=extras_require,
     version=versioneer.get_version(),