diff --git a/prod.Dockerfile b/prod.Dockerfile
index 0da5363a..96aa9d3c 100644
--- a/prod.Dockerfile
+++ b/prod.Dockerfile
@@ -1,4 +1,5 @@
-FROM ubuntu:jammy
+FROM nvidia/cuda:12.2.0-devel-ubuntu22.04
+# FROM ubuntu:jammy
 
 LABEL org.opencontainers.image.source https://github.com/khoj-ai/khoj