diff --git a/scripts/beaker/Dockerfile b/scripts/beaker/Dockerfile
index 9633f09..c6a2b5e 100644
--- a/scripts/beaker/Dockerfile
+++ b/scripts/beaker/Dockerfile
@@ -1,3 +1,20 @@
 FROM gcr.io/ai2-beaker-core/public/cqgl31u2ba5vrtuc91og:latest
 
-RUN apt update && apt-get install -y libaio-dev
\ No newline at end of file
+# Update the package list and install libaio-dev and gnupg2
+RUN apt update && apt-get install -y libaio-dev gnupg2
+
+# Add NVIDIA package repository keys
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub \
+    && apt-get -y update
+
+# Set up the NVIDIA CUDA repository
+RUN apt-get install -y software-properties-common \
+    && add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /" \
+    && apt-get update
+
+# Install CUDA toolkit and nvcc 12.1
+RUN apt-get install -y cuda-nvcc-12-1
+
+# Get flash attention setup
+RUN pip install flash-attn --no-build-isolation
\ No newline at end of file