diff --git a/entityreference_modal_create.module b/entityreference_modal_create.module
index fb0dee5..c864cda 100644
--- a/entityreference_modal_create.module
+++ b/entityreference_modal_create.module
@@ -56,9 +56,9 @@ function entityreference_modal_create_entity_type_build(array &$entity_types) {
   foreach ($entity_types as $typeId => $entity_type) {
     if(get_class($entity_type) == 'Drupal\Core\Entity\ContentEntityType') {
       $form_modes = \Drupal::service('entity_display.repository')->getFormModeOptions($typeId);
-      if($entity_type->getFormClass('edit')){
+      if($entity_type->getFormClass('default')){
         foreach ($form_modes as $modeKey => $label) {
-          $entity_type->setFormClass($modeKey, $entity_type->getFormClass('edit'));
+          $entity_type->setFormClass($modeKey, $entity_type->getFormClass('default'));
         }
       }
     }