diff --git a/unity-command/action.yaml b/unity-command/action.yaml
index 24e11ba..3e7d406 100644
--- a/unity-command/action.yaml
+++ b/unity-command/action.yaml
@@ -27,6 +27,10 @@ inputs:
   unityBuilder:
     description: "Whether or not to use the UnityBuilder instead of a Unity command."
     required: false
+  removeContainer:
+    description: "Remove the mock container after building."
+    required: false
+    default: "true"
 runs:
   using: "composite"
   steps:
@@ -40,7 +44,7 @@ runs:
     run: |
       CONTAINER=${{ steps.getContainer.outputs.container }}
 
-      #docker build -t ${{ inputs.imageTag }} --build-arg IMAGE=$CONTAINER ${{ github.action_path }}
+      docker build -t ${{ inputs.imageTag }} --build-arg IMAGE=$CONTAINER ${{ github.action_path }}
     shell: bash
   - name: "Run Unity command."
     uses: act/unity/unity@master
@@ -50,8 +54,8 @@ runs:
       password: ${{ inputs.password }}
       command: ${{ inputs.command }}
       unityBuilder: ${{ inputs.unityBuilder }}
-      image: docker://${{ steps.getContainer.outputs.container }}
   - name: "Remove temporary image."
+    if: ${{ inputs.removeContainer == "true" }}
     run: |
       docker image rm ${{ inputs.imageTag }}
     shell: bash