diff --git a/app/docker/components/container-restart-policy/container-restart-policy-controller.js b/app/docker/components/container-restart-policy/container-restart-policy-controller.js
index c9ece16d4..fd96a1596 100644
--- a/app/docker/components/container-restart-policy/container-restart-policy-controller.js
+++ b/app/docker/components/container-restart-policy/container-restart-policy-controller.js
@@ -1,36 +1,26 @@
 angular
-  .module('portainer.docker')
-  .controller('ContainerRestartPolicyController', [
-    function ContainerRestartPolicyController() {
-      var ctrl = this;
+.module('portainer.docker')
+.controller('ContainerRestartPolicyController', [function ContainerRestartPolicyController() {
+  var ctrl = this;
 
-      this.state = {
-        editMode :false,
-        editModel :{}
-      };
-      
+  this.state = {
+    editModel : {}
+  };
 
-      ctrl.toggleEdit = toggleEdit;
-      ctrl.save = save;
+  ctrl.save = save;
 
-      function toggleEdit() {
-        ctrl.state.editMode = true;
-        ctrl.state.editModel = {
-          name: ctrl.name,
-          maximumRetryCount: ctrl.maximumRetryCount
-        };
-      }
-
-      function save() {
-        if (ctrl.state.editModel.name === ctrl.name &&
-            ctrl.state.editModel.maximumRetryCount === ctrl.maximumRetryCount) {
-          ctrl.state.editMode = false;
-          return;
-        }
-        ctrl.updateRestartPolicy(ctrl.state.editModel)
-          .then(function onUpdateSucceed() {
-            ctrl.state.editMode = false;
-          });
-      }
+  function save() {
+    if (ctrl.state.editModel.name === ctrl.name && ctrl.state.editModel.maximumRetryCount === ctrl.maximumRetryCount) {
+      return;
     }
-  ]);
+    ctrl.updateRestartPolicy(ctrl.state.editModel);
+  }
+
+  this.$onInit = function() {
+    ctrl.state.editModel = {
+      name: ctrl.name,
+      maximumRetryCount: ctrl.maximumRetryCount
+    };
+  };
+}
+]);
diff --git a/app/docker/components/container-restart-policy/container-restart-policy.html b/app/docker/components/container-restart-policy/container-restart-policy.html
index c5352b6e4..41b87c2fc 100644
--- a/app/docker/components/container-restart-policy/container-restart-policy.html
+++ b/app/docker/components/container-restart-policy/container-restart-policy.html
@@ -1,23 +1,5 @@
 <div>
-  <table class="table table-bordered table-condensed" ng-if="!$ctrl.state.editMode">
-    <tr>
-      <td class="col-md-3">
-        <a href="" data-toggle="tooltip" title="Edit restart policy" ng-click="$ctrl.toggleEdit()">
-          <i class="fa fa-edit"></i>
-        </a>
-        <span>Name</span>
-      </td>
-      <td>{{$ctrl.name }}</td>
-    </tr>
-    <tr ng-if="$ctrl.name === 'on-failure'">
-      <td class="col-md-3">Maximum Retry Count</td>
-      <td>
-        {{ $ctrl.maximumRetryCount }}
-      </td>
-
-    </tr>
-  </table>
-  <table class="table table-bordered table-condensed" ng-if="$ctrl.state.editMode">
+  <table class="table table-bordered table-condensed">
     <tr>
       <td class="col-md-3">
         <span>Name</span>
@@ -31,15 +13,14 @@
         </select>
       </td>
       <td class="col-md-2">
-        <button class="btn btn-success" ng-click="$ctrl.save()">Save</button>
+        <button class="btn btn-sm btn-primary" ng-click="$ctrl.save()">Update</button>
       </td>
     </tr>
     <tr ng-if="$ctrl.state.editModel.name === 'on-failure'">
       <td class="col-md-3">Maximum Retry Count</td>
-      <td>
+      <td colspan="2">
         <input type="number" class="form-control" ng-model="$ctrl.state.editModel.maximumRetryCount" />
       </td>
-      <td class="col-md-2"></td>
     </tr>
   </table>
-</div>
\ No newline at end of file
+</div>
diff --git a/app/docker/views/containers/edit/container.html b/app/docker/views/containers/edit/container.html
index 92815f9c0..a251ae2d2 100644
--- a/app/docker/views/containers/edit/container.html
+++ b/app/docker/views/containers/edit/container.html
@@ -226,7 +226,7 @@
               <tr>
                 <td>Restart policies</td>
                 <td>
-                  <container-restart-policy 
+                  <container-restart-policy ng-if="container"
                     name="container.HostConfig.RestartPolicy.Name"
                     maximum-retry-count="container.HostConfig.RestartPolicy.MaximumRetryCount"
                     update-restart-policy="updateRestartPolicy(name, maximumRetryCount)">
diff --git a/app/docker/views/containers/edit/containerController.js b/app/docker/views/containers/edit/containerController.js
index e7b1a988b..f687b116c 100644
--- a/app/docker/views/containers/edit/containerController.js
+++ b/app/docker/views/containers/edit/containerController.js
@@ -323,6 +323,7 @@ function ($q, $scope, $state, $transition$, $filter, Commit, ContainerHelper, Co
         Name: restartPolicy,
         MaximumRetryCount: maximumRetryCount
       };
+      Notifications.success('Restart policy updated');
     }
 
     function notifyOnError(err) {