diff --git a/inline_entity_form.module b/inline_entity_form.module
index bb433cf..983cf6e 100644
--- a/inline_entity_form.module
+++ b/inline_entity_form.module
@@ -269,15 +269,8 @@ function inline_entity_form_close_all_forms($elements, FormStateInterface $form_
     $form_state->set(['inline_entity_form', $ief_id, 'form'], NULL);
     // Close the row forms.
     $entities = $form_state->get(['inline_entity_form', $ief_id, 'entities']);
-    if (!empty($entities)) {
-      if (isset($elements['#ief_row_delta']) && isset($entities[$elements['#ief_row_delta']])) {
-        $entities[$elements['#ief_row_delta']]['form'] = NULL;
-      }
-      else {
-        foreach ($entities as $key => $value) {
-          $entities[$key]['form'] = NULL;
-        }
-      }
+    foreach ($entities as $key => $value) {
+      $entities[$key]['form'] = NULL;
     }
     $form_state->set(['inline_entity_form', $ief_id, 'entities'], $entities);
   }
diff --git a/src/Plugin/Field/FieldWidget/InlineEntityFormComplex.php b/src/Plugin/Field/FieldWidget/InlineEntityFormComplex.php
index 38aaee8..e4615a1 100644
--- a/src/Plugin/Field/FieldWidget/InlineEntityFormComplex.php
+++ b/src/Plugin/Field/FieldWidget/InlineEntityFormComplex.php
@@ -626,8 +626,8 @@ class InlineEntityFormComplex extends InlineEntityFormBase implements ContainerF
     // Add submit handlers depending on operation.
     if ($element['#op'] == 'add') {
       static::addSubmitCallbacks($element['actions']['ief_add_save']);
+      $element['actions']['ief_add_save']['#submit'][] = [get_called_class(), 'closeForm'];
       $element['actions']['ief_add_cancel']['#submit'] = [
-        [get_called_class(), 'closeChildForms'],
         [get_called_class(), 'closeForm'],
         'inline_entity_form_cleanup_form_state',
       ];
@@ -637,7 +637,9 @@ class InlineEntityFormComplex extends InlineEntityFormBase implements ContainerF
       $element['actions']['ief_edit_cancel']['#ief_row_delta'] = $element['#ief_row_delta'];
 
       static::addSubmitCallbacks($element['actions']['ief_edit_save']);
+      $element['actions']['ief_edit_save']['#submit'][] = [get_called_class(), 'closeChildForms'];
       $element['actions']['ief_edit_save']['#submit'][] = [get_called_class(), 'submitCloseRow'];
+
       $element['actions']['ief_edit_cancel']['#submit'] = [
         [get_called_class(), 'closeChildForms'],
         [get_called_class(), 'submitCloseRow'],
@@ -876,7 +878,6 @@ class InlineEntityFormComplex extends InlineEntityFormBase implements ContainerF
   public static function addSubmitCallbacks(&$element) {
     $element['#submit'] = [
       ['\Drupal\inline_entity_form\ElementSubmit', 'trigger'],
-      ['\Drupal\inline_entity_form\Plugin\Field\FieldWidget\InlineEntityFormComplex', 'closeForm'],
     ];
     $element['#ief_submit_trigger']  = TRUE;
   }
@@ -893,16 +894,14 @@ class InlineEntityFormComplex extends InlineEntityFormBase implements ContainerF
    *   The form state of the parent form.
    */
   public static function closeChildForms($form, FormStateInterface &$form_state) {
-    $element = inline_entity_form_get_element($form, $form_state);
+    // Find the closest inline_entity_form element up the tree
     $triggering_element = $form_state->getTriggeringElement();
-    if (!preg_match('~^ief-add-cancel~i', $triggering_element['#name'])) {
-      $array_parents = array_reverse($triggering_element['#array_parents']);
-      while (current($array_parents) != 'inline_entity_form') {
-        array_shift($array_parents);
-      }
-      $array_parents = array_reverse($array_parents);
-      $element = NestedArray::getValue($form, $array_parents) ?: $element;
-      inline_entity_form_close_all_forms($element, $form_state);
+    // Remove the action and the actions container.
+    $array_parents = array_slice($triggering_element['#array_parents'], 0, -2);
+    $inline_entity_form = NestedArray::getValue($form, $array_parents);
+    foreach (Element::children($inline_entity_form) as $key) {
+      inline_entity_form_close_all_forms($inline_entity_form[$key], $form_state);
     }
   }
+
 }
diff --git a/src/Tests/ComplexSimpleWidgetTest.php b/src/Tests/ComplexSimpleWidgetTest.php
index 92e6f74..550112f 100644
--- a/src/Tests/ComplexSimpleWidgetTest.php
+++ b/src/Tests/ComplexSimpleWidgetTest.php
@@ -133,7 +133,7 @@ class ComplexSimpleWidgetTest extends InlineEntityFormTestBase {
     $this->drupalPostAjaxForm(NULL, [], $this->getButtonName('//input[@type="submit" and @value="Cancel" and @data-drupal-selector="edit-ief-complex-outer-form-inline-entity-form-entities-0-form-actions-ief-edit-cancel"]'));
     $this->assertFieldByName('ief_complex_outer[form][inline_entity_form][entities][1][form][title][0][value]', '', 'Second inline edit form remained open when the first was closed.');
 
-    // Second and third assertion. In both cases, we open and edit form and the
+    // Second and third assertion. In both cases, we open an edit form and the
     // create form. We then try hitting Cancel on the edit and create in turn
     // making sure the other remains open.
     $create_clicks = array(TRUE, FALSE);