diff --git a/dev-shell-e2e-test.sh b/dev-shell-e2e-test.sh
index 39d2a5e3..1d1d1f59 100755
--- a/dev-shell-e2e-test.sh
+++ b/dev-shell-e2e-test.sh
@@ -14,4 +14,4 @@ down() {
 trap down EXIT
 
 mkdir -p $PWD/.device/sysroot
-docker compose --env-file=${docker_path}/.env.dev -f ${docker_path}/docker-compose.yml run -e DEV_USER=$(id -u) -e DEV_GROUP=$(id -g) -e BASE_TARGET_VERSION=${BASE_TARGET_VERSION} -e USER_TOKEN=${USER_TOKEN} aklite-e2e-test "$@"
+docker compose --env-file=${docker_path}/.env.dev -f ${docker_path}/docker-compose.yml run -e DEV_USER=$(id -u) -e DEV_GROUP=$(id -g) -e BASE_TARGET_VERSION=${BASE_TARGET_VERSION} -e USER_TOKEN=${USER_TOKEN} -e TAG=${TAG} aklite-e2e-test "$@"
diff --git a/e2e-test.py b/e2e-test.py
index 466ed51f..3dfb053b 100644
--- a/e2e-test.py
+++ b/e2e-test.py
@@ -59,6 +59,11 @@ class ReturnCodes:
     logger.error("BASE_TARGET_VERSION environment variable not set")
     sys.exit()
 
+tag = os.getenv("TAG")
+if not tag:
+    logger.error("TAG environment variable not set")
+    sys.exit()
+
 base_target_version = int(base_target_version)
 logger.info(f"Base target version: {base_target_version}")
 
@@ -118,7 +123,7 @@ def __str__(self):
 def register_if_required():
     if not os.path.exists("/var/sota/client.pem"):
         user_token = os.getenv("USER_TOKEN")
-        cmd = f'DEVICE_FACTORY={factory_name} lmp-device-register --api-token "{user_token}" --start-daemon 0 --tags main'
+        cmd = f'DEVICE_FACTORY={factory_name} lmp-device-register --api-token "{user_token}" --start-daemon 0 --tags {tag}'
         logger.info(f"Registering device...")
         output = os.popen(cmd).read().strip()
         logger.info(output)
@@ -224,7 +229,7 @@ def write_settings(apps=None, prune=True):
     content = \
 f"""
 [pacman]
-tags = "main"
+tags = "{tag}"
 """
     if apps is not None:
         apps_str = ",".join(apps)