diff --git a/entity_forms_in_blocks/entity_forms_in_blocks.module b/entity_forms_in_blocks/entity_forms_in_blocks.module
index bfa098f..1ad4c34 100644
--- a/entity_forms_in_blocks/entity_forms_in_blocks.module
+++ b/entity_forms_in_blocks/entity_forms_in_blocks.module
@@ -1,19 +1,22 @@
 <?php
 /**
- * Implements hook_entity_form_display_alter().
+ * Implements hook_entity_type_alter().
  */
-function entity_forms_in_blocks_entity_form_display_alter(&$form_display, $context) {
-  // Load the right entity form display. Works for any entity / bundle.
-  global $current_context;
-  $id =implode('.', $current_context);
-  $storage = \Drupal::entityManager()->getStorage('entity_form_display');
-  $change_display = $storage->load($id);
-  unset($current_context);
-  // If form mode is activated, replace the given one with ours.
-  if ($change_display) {
-    $form_display = $change_display;
-  }
+function entity_forms_in_blocks_entity_type_alter(array &$entity_types) {
+    $form_modes = \Drupal::service('entity_display.repository')->getAllFormModes();
+    foreach ($form_modes as $entity_type => $display_modes) {
+        /* @var \Drupal\Core\Entity\EntityTypeInterface $entity */
+        $entity = $entity_types[$entity_type];
+        foreach ($display_modes as $machine_name => $form_display) {
+            if ($machine_name != 'register') {
+                // Get the correct canonical path to add operation.
+                $path = $entity->getLinkTemplate('canonical') . "/$machine_name";
+                $default_handler_class = $entity->getHandlerClasses()['form']['default'];
+                $entity->setFormClass($machine_name, $default_handler_class)
+                    ->setLinkTemplate($machine_name, $path);
+            }
+        }
+    }
 }
 
 
-
diff --git a/entity_forms_in_blocks/src/Plugin/Block/EntityFormBlock.php b/entity_forms_in_blocks/src/Plugin/Block/EntityFormBlock.php
index 35eaf34..a33a39a 100644
--- a/entity_forms_in_blocks/src/Plugin/Block/EntityFormBlock.php
+++ b/entity_forms_in_blocks/src/Plugin/Block/EntityFormBlock.php
@@ -280,7 +280,7 @@ class EntityFormBlock extends BlockBase implements ContainerFactoryPluginInterfa
     }
     $entity = $this->entityManager->getStorage($definition->get('id'))
       ->create($values);
-    return $this->entityFormBuilder->getForm($entity);
+      return $this->entityFormBuilder->getForm($entity, $this->configuration['form_mode']);
   }