diff --git a/deployment/Dockerfile.prediction.lambda b/deployment/Dockerfile.prediction.lambda
index 8a759fd..a011d55 100644
--- a/deployment/Dockerfile.prediction.lambda
+++ b/deployment/Dockerfile.prediction.lambda
@@ -4,6 +4,10 @@ FROM public.ecr.aws/lambda/python:3.10
 WORKDIR ${LAMBDA_TASK_ROOT}
 ENV PYTHONPATH "${PYTHONPATH}:${LAMBDA_TASK_ROOT}"
 
+# Environment variables
+ARG RUNTIME_ENVIRONMENT
+ENV RUNTIME_ENVIRONMENT=${RUNTIME_ENVIRONMENT}
+
 # Install necessary build tools - required to test locally
 RUN yum install -y gcc python3-devel
 
@@ -16,6 +20,6 @@ COPY modules/ml-pipeline/src/pipeline ./pipeline
 # Copy the handler
 COPY deployment/handlers/prediction_app.py prediction_app.py
 # Get the model
-# RUN dev pull -r ${RUNTIME_ENVIRONMENT}
+# RUN dvc pull -r ${RUNTIME_ENVIRONMENT}
 
 CMD [ "prediction_handler.handler" ]