diff --git a/docker/Dockerfile-beta b/docker/Dockerfile-beta
index 04f15f21c7..6b7a84b70e 100644
--- a/docker/Dockerfile-beta
+++ b/docker/Dockerfile-beta
@@ -4,7 +4,9 @@ RUN apt update && apt install -y python3.10-venv git-lfs
 
 RUN python3 -m pip install --upgrade pip &&  \
     pip install packaging && \
-    pip install -U torch==2.2.2
+    pip uninstall torch-tensorrt && \
+    pip install -U torch==2.2.2 && \
+    pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
 
 RUN groupadd axolotl &&  \
     useradd -m -g axolotl -s /bin/bash axolotl &&  \
@@ -20,8 +22,7 @@ ENV PATH="/home/axolotl/venv/axolotl/bin:$PATH"
 
 RUN echo "source /home/axolotl/venv/axolotl/bin/activate" >> /home/axolotl/.bashrc
 
-RUN git lfs install --skip-repo && \
-    pip3 install awscli
+RUN git lfs install --skip-repo
 
 WORKDIR /workspace
 
diff --git a/setup.py b/setup.py
index fbca5a360e..3d221cf5e9 100644
--- a/setup.py
+++ b/setup.py
@@ -68,7 +68,7 @@ def parse_requirements():
     dependency_links=dependency_links,
     extras_require={
         "flash-attn": [
-            "flash-attn==2.5.5",
+            "flash-attn>=2.4.2",
         ],
         "fused-dense-lib": [
             "fused-dense-lib  @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",