diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 55afbc8..b831a5d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -257,10 +257,10 @@ jobs:
         with:
           context: .
           build-args: |
-            BASE_IMAGE=nvcr.io/nvidia/pytorch:23.06-py3
+            BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu22.04
             GPU_ENABLED=true
           # BASE_IMAGE=nvcr.io/nvidia/pytorch:23.06-py3
-          # BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu20.04
+          # BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu22.04
           push: ${{ github.event_name != 'pull_request' }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
diff --git a/Dockerfile b/Dockerfile
index 0fc31d2..6a7295e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,9 +1,11 @@
 ARG BASE_IMAGE=python:3.11
-# ARG BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu20.04
+# ARG BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu22.04
 # 2.7GB cf. https://ngc.nvidia.com/catalog/containers/nvidia:cuda
 # ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:23.06-py3
 # 8.5GB cf. https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
 
+# CUDA image: https://github.com/oobabooga/text-generation-webui/blob/main/docker/Dockerfile
+
 FROM ${BASE_IMAGE}
 
 ARG DEBIAN_FRONTEND=noninteractive
@@ -12,14 +14,14 @@ ENV LIBRECHAT_WORKERS=1
 
 # CUDA image required to install python
 RUN apt-get update && \
-    apt-get install -y software-properties-common wget unzip && \
+    apt-get install -y software-properties-common git vim build-essential python3-dev wget unzip && \
     # add-apt-repository ppa:deadsnakes/ppa && \
     # apt-get install -y python3.11 && \
     # ln -s /usr/bin/python3.11 /usr/bin/python && \
     # wget https://bootstrap.pypa.io/get-pip.py && \
     # python get-pip.py && \
     # rm get-pip.py && \
-    pip install --upgrade pip
+    pip3 install --upgrade pip
 
 
 
@@ -37,7 +39,7 @@ WORKDIR /app
 
 # Pre-install requirements to use cache when re-building
 ADD scripts/requirements.txt .
-RUN pip install -r requirements.txt && \
+RUN pip3 install -r requirements.txt && \
     rm requirements.txt
 
 ADD . .
diff --git a/src/libre_chat/ui.py b/src/libre_chat/ui.py
index c6d59a0..616211d 100644
--- a/src/libre_chat/ui.py
+++ b/src/libre_chat/ui.py
@@ -51,7 +51,7 @@ def stream(
         input_text: str, memory: Any, instructions: str, temperature: float, max_new_tokens: int
     ) -> Iterator[Tuple[Any, str]]:
         # Create a Queue
-        q: Queue[Any] = Queue()
+        q: "Queue[Any]" = Queue()
         job_done = object()
 
         # Create a function to call - this will run in a thread
@@ -240,7 +240,7 @@ def clear_chat() -> List[str]:
 class StreamGradioCallback(BaseCallbackHandler):
     """Callback handler for streaming LLM responses to a queue."""
 
-    def __init__(self, q: Queue[Any]):
+    def __init__(self, q: "Queue[Any]"):
         self.q = q
 
     def on_llm_new_token(self, token: str, **kwargs: Any) -> None: