diff --git a/contrib/images/babylond/Dockerfile b/contrib/images/babylond/Dockerfile
index c243ae5e4..ce6a6564a 100644
--- a/contrib/images/babylond/Dockerfile
+++ b/contrib/images/babylond/Dockerfile
@@ -1,14 +1,20 @@
-FROM --platform=$BUILDPLATFORM golang:1.21 AS build-env
+FROM golang:1.21 AS build-env
 
+# Version to build. Default is empty
+ARG VERSION
 ARG BUILD_TAGS=""
 ARG LEDGER_ENABLED="false"
 ARG COSMOS_BUILD_OPTIONS=""
 
-ARG TARGETOS TARGETARCH
-
 WORKDIR /go/src/github.com/babylonlabs-io/babylon
 COPY ./ /go/src/github.com/babylonlabs-io/babylon/
 
+# Handle if version is set
+RUN if [ -n "${VERSION}" ]; then \
+        git fetch origin tag ${VERSION} --no-tags; \
+        git checkout -f ${VERSION}; \
+    fi
+
 # Cache mod dependencies
 RUN go mod download