diff --git a/.github/workflows/Deploy.yml b/.github/workflows/Deploy.yml
index f97bddc..67d8cb6 100644
--- a/.github/workflows/Deploy.yml
+++ b/.github/workflows/Deploy.yml
@@ -91,14 +91,16 @@ jobs:
       - name: Build and push Docker image
         uses: docker/build-push-action@v3
         with:
-          context: ./model_data/simulation_system
-          file: ./model_data/simulation_system/Dockerfiles/Dockerfile.prediction.lambda
+          context: ./modules/ml-pipeline/src/pipeline
+          file: ./deployment/Dockerfile.prediction.lambda
           push: true
           tags: ${{ steps.set_ecr_credentials.outputs.ecr_uri }}:${{ github.sha }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
           platform: linux/amd64
           provenance: false
+          build-args: |
+            RUNTIME_ENVIRONMENT=${{ steps.set_runtime_environment.outputs.runtime_environment }}
 
       - name: Deploy to AWS Lambda via Serverless
         env:
@@ -112,4 +114,4 @@ jobs:
           STACK_NAME: ${{ steps.set_stack_name.outputs.stack_name }}
         run: |
           # Deploy to AWS Lambda via Serverless
-          sls deploy --config serverless.yml --stage ${{ steps.set_runtime_environment.outputs.runtime_environment }} --verbose
+          sls deploy --config deployment/serverless.yml --stage ${{ steps.set_runtime_environment.outputs.runtime_environment }} --verbose