diff --git a/.gitignore b/.gitignore
index d8a9ff3..a435248 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,4 @@ mlruns/
 output/
 outputs/
 *.egg-info/
+.ipynb_checkpoints/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f58c5ae..9374303 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -17,7 +17,7 @@ repos:
           - --line-length=120
           - --safe
           - --target-version=py310
-        stages: [commit]
+        stages: [pre-commit]
   - repo: https://github.com/pycqa/isort
     rev: 5.12.0
     hooks:
diff --git a/Dockerfile b/Dockerfile
index 7b290b9..d40e24c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,22 +1,31 @@
 # https://hub.docker.com/layers/winglian/axolotl/main-20241104-py3.11-cu121-2.3.1/images/sha256-790297fa1d71f8f1590c73ca4505ca39fe7dfa2886b2b6862199a6df679bf8e4?context=explore
 FROM winglian/axolotl@sha256:cffbcc4993e80301a8918062f8136a6ac402877fd6c29f1168be563e543aee4d
+SHELL ["/bin/bash", "-c"]
 USER root
+
+# Install torch
 COPY requirements.txt /tmp/
 RUN pip install -U pip wheel setuptools && \
     pip uninstall -y axolotl && \
     pip install --no-cache-dir -U -r /tmp/requirements.txt
+
+# Install axolotl
 RUN mkdir -p /packages && \
     cd /packages && \
     git clone https://github.com/truefoundry/axolotl && \
     cd axolotl/ && \
-    git checkout e16f637d079ef5d56321a240ef0547a50c37b708
-RUN cd /packages/axolotl/ && \
+    git checkout e16f637d079ef5d56321a240ef0547a50c37b708 && \
+    cd /packages/axolotl/ && \
     MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,fused-dense-lib,optimizers,lion-pytorch,galore] && \
-    pip install --no-cache-dir -U -r /tmp/requirements.txt && \
     rm -rf /root/.cache/pip
+
+# Install axolotl_truefoundry plugin with our requirements overrides over axolotl
 COPY plugins/axolotl_truefoundry /packages/axolotl_truefoundry
 RUN cd /packages/axolotl_truefoundry/ && \
-    pip install --no-cache-dir -e . && \
+    pip install --no-cache-dir -U -r /tmp/requirements.txt -e . && \
     rm -rf /root/.cache/pip
+
 WORKDIR /app
+
+# Add source code for finetuning
 COPY . /app
diff --git a/Dockerfile-notebook b/Dockerfile-notebook
index 85505bd..d43fc46 100644
--- a/Dockerfile-notebook
+++ b/Dockerfile-notebook
@@ -2,6 +2,8 @@ FROM tfy.jfrog.io/tfy-images/jupyter:0.3.4-cu121-py3.11.10-sudo
 SHELL ["/bin/bash", "-c"]
 ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
 ENV DEBIAN_FRONTEND=noninteractive
+
+# upgrade libnccl
 USER root
 RUN apt update && \
     apt install -y --no-install-recommends git curl wget && \
@@ -9,25 +11,33 @@ RUN apt update && \
     dpkg -i /tmp/cuda-keyring_1.1-1_all.deb && \
     apt update && \
     apt install -y --no-install-recommends --allow-change-held-packages libnccl2 libnccl-dev
+
+# Install torch
 USER jovyan
 COPY requirements.txt notebook-requirements.txt /tmp/llm-finetune/
 RUN pip install -U pip wheel setuptools && \
     pip uninstall -y axolotl && \
     pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt
+
+# Setup editable packages
 USER root
 RUN mkdir -p /packages && \
     chown -R jovyan:users /packages
+
+# Install axolotl
 USER jovyan
-RUN pip install -U pip setuptools wheel && \
-    cd /packages && \
+RUN cd /packages && \
     git clone https://github.com/truefoundry/axolotl && \
     cd axolotl/ && \
-    git checkout e16f637d079ef5d56321a240ef0547a50c37b708
-RUN cd /packages/axolotl/ && \
-    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,fused-dense-lib,optimizers,lion-pytorch,galore] && \
-    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt
+    git checkout e16f637d079ef5d56321a240ef0547a50c37b708 && \
+    cd /packages/axolotl/ && \
+    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,fused-dense-lib,optimizers,lion-pytorch,galore]
+
+# Install axolotl_truefoundry plugin with our requirements overrides over axolotl
 COPY --chown=jovyan:users plugins/axolotl_truefoundry /packages/axolotl_truefoundry
 RUN cd /packages/axolotl_truefoundry/ && \
-    pip install --no-cache-dir -e .
+    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt -e .
+
+# Add source code for finetuning
 COPY --chown=jovyan:users . /tmp_home/jovyan/llm-finetune/
 ENV JUPYTER_APP_LAUNCHER_PATH=/home/jovyan/llm-finetune/.jp_app_launcher/
diff --git a/plugins/axolotl_truefoundry/pyproject.toml b/plugins/axolotl_truefoundry/pyproject.toml
index ac0209b..d55dd19 100644
--- a/plugins/axolotl_truefoundry/pyproject.toml
+++ b/plugins/axolotl_truefoundry/pyproject.toml
@@ -9,9 +9,9 @@ description = "TrueFoundry plugin for Axolotl"
 requires-python = ">=3.8.1,<4.0"
 dependencies = [
     "transformers>=4.0.0,<5",
-    "truefoundry==0.4.4",
+    "truefoundry>=0.4.4,<0.5.0",
     "pynvml>=11.0.0,<12",
-    "torch>=2.3.0,<3",
+    "torch>=2.3.0,<2.4.0",
     "pydantic>=2.0.0,<3",
 ]
 
diff --git a/requirements.txt b/requirements.txt
index 888fefc..40fa7e4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,5 +7,5 @@ rich>=13.0.0,<14
 snowflake-connector-python[pandas]==3.12.3
 torch==2.3.1+cu121
 torchao==0.6.1+cu121
-truefoundry==0.4.4
+truefoundry>=0.4.4,<0.5.0
 unsloth @ git+https://github.com/unslothai/unsloth@38663b01f5dd0e610b12475bd95b144303cff539