diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml
index eac3d04b3..7b395f46c 100644
--- a/.github/workflows/format.yml
+++ b/.github/workflows/format.yml
@@ -32,7 +32,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install tox pre-commit
+          python -m pip install tox
       - name: Check formatting
         run: tox -e fmt
       - name: Run pylint
diff --git a/.github/workflows/image.yaml b/.github/workflows/image.yaml
new file mode 100644
index 000000000..2af2f086a
--- /dev/null
+++ b/.github/workflows/image.yaml
@@ -0,0 +1,25 @@
+name: Image
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.11
+      - name: Install dependencies
+        run: |
+          python -m pip install tox
+      - name: Build wheel
+        run: |
+          tox -e build
+      - name: Build image
+        run: |
+          docker build -t fms-hf-tuning:dev . -f build/Dockerfile
\ No newline at end of file
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 64720779f..24c3be205 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -15,6 +15,8 @@ jobs:
             tox: "py39"
           - setup: "3.10"
             tox: "py310"
+          - setup: "3.11"
+            tox: "py311"
     steps:
       - uses: actions/checkout@v3
       - name: Install dependencies
@@ -22,7 +24,4 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install tox build
       - name: Run unit tests
-        run: tox -e py
-      # TODO: switch to tox
-      - name: Build wheel
-        run: python -m build .
\ No newline at end of file
+        run: tox -e py
\ No newline at end of file
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 2017e2670..9a79627a4 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -38,9 +38,9 @@ Before sending pull requests, make sure your changes pass formatting, linting an
 #### Dependencies
 If additional new Python module dependencies are required, think about where to put them:
 
-- If they're required for fms-hf-tuning, then append them to the end of the requirements.txt file.
-- If they're optional dependencies for additional functionality, then put them in the pyproject.toml file like were done for [flash-attn](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L26) or [aim](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L27).
-- If it's an additional dependency for development, then add it to the [dev](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L25) dependencies.
+- If they're required for fms-hf-tuning, then append them to the [dependencies](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L27) in the pyproject.toml.
+- If they're optional dependencies for additional functionality, then put them in the pyproject.toml file like were done for [flash-attn](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L43) or [aim](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L44).
+- If it's an additional dependency for development, then add it to the [dev](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/pyproject.toml#L42) dependencies.
 
 #### Code Review
 
diff --git a/build/Dockerfile b/build/Dockerfile
index 2b0935084..7909c6857 100644
--- a/build/Dockerfile
+++ b/build/Dockerfile
@@ -12,6 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+FROM registry.access.redhat.com/ubi9/python-311 as build
+
+USER root
+RUN python -m pip install --upgrade pip && \
+    python -m pip install build
+COPY tuning tuning
+COPY pyproject.toml .
+RUN --mount=source=.git,target=.git,type=bind \
+    python -m build --wheel --outdir /tmp && \
+    ls /tmp/*.whl >/tmp/bdist_name
+
+
 FROM registry.access.redhat.com/ubi9/ubi AS release
 
 ARG CUDA_VERSION=11.8.0
@@ -94,22 +106,18 @@ RUN dnf install -y python3.11 git && \
     ln -s /usr/bin/python3.11 /bin/python && \
     python -m ensurepip --upgrade
 
-RUN mkdir /app
-
 WORKDIR /tmp
-RUN python -m pip install --upgrade pip
-
-# TODO Move to installing wheel once we have proper releases setup instead of cloning the repo
-RUN git clone https://github.com/foundation-model-stack/fms-hf-tuning.git && \
-    cd fms-hf-tuning && \
-    python -m pip install ".[dev]" && \
-    python -m pip install ".[flash-attn]" && \
-    python -m pip install ".[aim]" && \
-    python -m pip install -U datasets
+COPY --from=build /tmp/*.whl /tmp/bdist_name /tmp/
+# ami and flash-attn depends on dev
+RUN --mount=type=cache,target=/root/.cache/pip \
+    python -m pip install "$(head bdist_name)[dev]" && \
+    python -m pip install "$(head bdist_name)[aim]" && \
+    python -m pip install "$(head bdist_name)[flash-attn]"
 
 RUN mkdir -p /licenses
 COPY LICENSE /licenses/
 
+RUN mkdir /app
 # Copy scripts and default configs
 COPY build/launch_training.py build/accelerate_launch.py fixtures/accelerate_fsdp_defaults.yaml /app/
 COPY build/utils.py /app/build/
diff --git a/pyproject.toml b/pyproject.toml
index 24847abf1..6fc727eed 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ requires = [
 
 [project]
 name = "fms-hf-tuning"
-dynamic = ["version", "dependencies"]
+dynamic = ["version"]
 description = "FMS HF Tuning"
 authors = [
   {name = "Sukriti Sharma", email = "sukriti.sharma4@ibm.com"},
@@ -50,6 +50,7 @@ namespaces = false
 [tool.setuptools_scm]
 version_file = "tuning/_version.py"
 
+
 [project.urls]
 Homepage = "https://github.com/foundation-model-stack/fms-hf-tuning"
 Repository = "https://github.com/foundation-model-stack/fms-hf-tuning"
diff --git a/setup.py b/setup.py
deleted file mode 100644
index d3cb2fd03..000000000
--- a/setup.py
+++ /dev/null
@@ -1,18 +0,0 @@
-# Copyright The FMS HF Tuning Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Third Party
-from setuptools import find_packages, setup
-
-setup(name="fms-hf-tuning", packages=find_packages())
diff --git a/tox.ini b/tox.ini
index 26eea355a..14dd1d715 100644
--- a/tox.ini
+++ b/tox.ini
@@ -10,14 +10,17 @@ commands =
 
 [testenv:fmt]
 description = format with pre-commit
+deps =
+    pre-commit
 commands = ./scripts/fmt.sh
 allowlist_externals = ./scripts/fmt.sh
 
 [testenv:lint]
 description = lint with pylint
-deps = pylint>=2.16.2,<=3.1.0
-        pytest
-        .[dev]
+deps =
+    pylint>=2.16.2,<=3.1.0
+    pytest
+    .[dev]
 commands = pylint tuning scripts/*.py build/*.py tests
 allowlist_externals = pylint