diff --git a/action.yml b/action.yml
index bffb1cba..025e7966 100644
--- a/action.yml
+++ b/action.yml
@@ -37,6 +37,11 @@ inputs:
   command:
     description: 'The command used by ECS to start the container image'
     required: false
+  task-role-arn:
+    description: 'The ARN of the IAM role that the ECS container will assume'
+    required: false
+  execution-role-arn:
+    description: 'The ARN of the IAM role that the ECS task execution role will assume'
   env-files:
     description: 'S3 object arns to set env variables onto the container. You can specify multiple files with multi-line YAML strings.'
     required: false
diff --git a/index.js b/index.js
index f4cb591d..ca18ca12 100644
--- a/index.js
+++ b/index.js
@@ -45,6 +45,9 @@ async function run() {
     const logConfigurationOptions = core.getInput("log-configuration-options", { required: false });
     const dockerLabels = core.getInput('docker-labels', { required: false });
     const command = core.getInput('command', { required: false });
+    const taskRoleArn = core.getInput('task-role-arn', { required: false });
+    const executionRoleArn = core.getInput('execution-role-arn', { required: false });
+
 
     //New inputs to fetch task definition 
     const taskDefinitionArn = core.getInput('task-definition-arn', { required: false }) || undefined;
@@ -109,6 +112,7 @@ async function run() {
     }
     containerDef.image = imageURI;
 
+
     if (command) {
       containerDef.command = command.split(' ')
     }
@@ -245,6 +249,14 @@ async function run() {
       })
     }
 
+    if (taskRoleArn) {
+      taskDefContents.taskRoleArn = taskRoleArn;
+    }
+
+    if (executionRoleArn) {
+      taskDefContents.executionRoleArn = executionRoleArn;
+    }
+
     // Write out a new task definition file
     var updatedTaskDefFile = tmp.fileSync({
       tmpdir: process.env.RUNNER_TEMP,