diff --git a/Dockerfile b/Dockerfile
index cdca38b..2fd6ddf 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,5 @@
-FROM cccs/assemblyline-v4-service-base:latest
+ARG branch=latest
+FROM cccs/assemblyline-v4-service-base:$branch
 
 ENV SERVICE_PATH avclass_.AVclass
 
diff --git a/pipelines/azure-build.yaml b/pipelines/azure-build.yaml
index 31d70f3..3059e6c 100644
--- a/pipelines/azure-build.yaml
+++ b/pipelines/azure-build.yaml
@@ -25,7 +25,7 @@ stages:
         set -xv  # Echo commands before they are run
         export TAG=${BUILD_SOURCEBRANCH#"refs/tags/v"}
         if [[ "$TAG" == *stable* ]]; then export BUILD_TYPE=stable; else export BUILD_TYPE=latest; fi
-        docker build --build-arg version=$TAG -t cccs/assemblyline-service-avclass:$TAG -t cccs/assemblyline-service-avclass:$BUILD_TYPE .
+        docker build --build-arg version=$TAG --build-arg branch=$BUILD_TYPE -t cccs/assemblyline-service-avclass:$TAG -t cccs/assemblyline-service-avclass:$BUILD_TYPE .
       displayName: Build containers
     - script: |
         docker run -v `pwd`/test/:/opt/al_service/test/ cccs/assemblyline-service-avclass:latest bash -c 'pip install -U -r test/requirements.txt; pytest'