diff --git a/.github/workflows/build-and_publish-prebuilts-whl.yml b/.github/workflows/build-and_publish-prebuilts-whl.yml
index f4df772..9848bee 100644
--- a/.github/workflows/build-and_publish-prebuilts-whl.yml
+++ b/.github/workflows/build-and_publish-prebuilts-whl.yml
@@ -41,14 +41,8 @@ jobs:
 
       # Step 4: Create a container and run the script
       - name: Run Docker container and generate wheels
-        env:
-          ROCM_VERSION: ${{ inputs.rocm_version }}
-          TORCH_VERSION: ${{ inputs.torch_version }}
         run: |
-          docker create --name prebuilt-container \
-            -e ROCM_ARCH=${ROCM_ARCH} \
-            -e TORCH_VERSION=${TORCH_VERSION} \
-            prebuilt-wheels
+          docker create --name prebuilt-container prebuilt-wheels
           docker start -a prebuilt-container
 
       # Step 5: Copy bitsandbytes wheel artifact to host
diff --git a/prebuilts/Dockerfile b/prebuilts/Dockerfile
index dba672a..2829975 100644
--- a/prebuilts/Dockerfile
+++ b/prebuilts/Dockerfile
@@ -1,3 +1,8 @@
+ARG ROCM_VERSION
+ARG TORCH_VERSION
+ENV ROCM_VERSION=${ROCM_VERSION}
+ENV TORCH_VERSION=${TORCH_VERSION}
+
 FROM rocm/dev-ubuntu-22.04:${ROCM_VERSION}
 
 ENV DEBIAN_FRONTEND=noninteractive \
@@ -6,7 +11,6 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 # For bitsandbytes
 ENV BNB_GPU_TARGETS="gfx803;gfx900;gfx906;gfx908;gfx90a;gfx1010;gfx1030;gfx1100;gfx1101;gfx1102"
-
 # For LLAMA
 ENV GPU_TARGETS="gfx803 gfx900 gfx906 gfx908 gfx90a gfx1010 gfx1030 gfx1100 gfx1101 gfx1102"