diff --git a/Dockerfile b/Dockerfile
index 317b9b09..ce25dd94 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,11 +1,17 @@
-FROM nvcr.io/nvidia/tensorrt:24.02-py3
+FROM nvidia/cuda:11.8.0-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()'
+# RUN python -c 'from rembg.bg import download_models; download_models()'
 
 EXPOSE 7000
 ENTRYPOINT ["rembg"]
diff --git a/setup.py b/setup.py
index 3332664e..22c8a516 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@
         "twine",
         "wheel",
     ],
-    "gpu": [""],
+    "gpu": ["onnxruntime-gpu"],
     "cli": [
         "aiohttp",
         "asyncer",