diff --git a/benchmarks/tract-onnx-image-classification/Dockerfile b/benchmarks/tract-onnx-image-classification/Dockerfile
index 3f50b95..cecfcd8 100644
--- a/benchmarks/tract-onnx-image-classification/Dockerfile
+++ b/benchmarks/tract-onnx-image-classification/Dockerfile
@@ -3,10 +3,11 @@ RUN rustup target add wasm32-wasi
 WORKDIR /usr/src
 ADD rust-benchmark rust-benchmark
 ADD assets assets
+ADD setup.sh setup.sh
+RUN ./setup.sh
 WORKDIR /usr/src/rust-benchmark
 ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse
 RUN cargo build --release --target wasm32-wasi
-RUN ls target/wasm32-wasi/release
 RUN mkdir /benchmark
 RUN cp target/wasm32-wasi/release/benchmark.wasm /benchmark/
 # We output the Wasm file to the `/benchmark` directory, where the client
diff --git a/benchmarks/tract-onnx-image-classification/assets/.gitignore b/benchmarks/tract-onnx-image-classification/assets/.gitignore
new file mode 100644
index 0000000..b722e9e
--- /dev/null
+++ b/benchmarks/tract-onnx-image-classification/assets/.gitignore
@@ -0,0 +1 @@
+!.gitignore
\ No newline at end of file
diff --git a/benchmarks/tract-onnx-image-classification/rust-benchmark/Cargo.toml b/benchmarks/tract-onnx-image-classification/rust-benchmark/Cargo.toml
index 89ccb7d..331672e 100644
--- a/benchmarks/tract-onnx-image-classification/rust-benchmark/Cargo.toml
+++ b/benchmarks/tract-onnx-image-classification/rust-benchmark/Cargo.toml
@@ -13,4 +13,4 @@ anyhow = "1.0"
 bytes = "1.5.0"
 prost = "0.11.0"
 image = { version = "0.25.1", features = ["png"], default-features = false }
-tract-onnx = { git = "https://github.com/sonos/tract", version = "=0.21.4-pre" }
+tract-onnx = "0.21.4"
diff --git a/benchmarks/tract-onnx-image-classification/setup.sh b/benchmarks/tract-onnx-image-classification/setup.sh
index 3c985dd..3634b29 100755
--- a/benchmarks/tract-onnx-image-classification/setup.sh
+++ b/benchmarks/tract-onnx-image-classification/setup.sh
@@ -5,6 +5,8 @@ set -euxo pipefail
 MODEL_URL="https://www.github.com/onnx/models/raw/main/validated/vision/classification/mobilenet/model/mobilenetv2-7.onnx?download=:"
 MODEL_FILE="assets/mobilenetv2-7.onnx"
 
+mkdir -p assets
+
 if [ ! -f $MODEL_FILE ]; then
 	echo "Downloading model to $MODEL_FILE"