diff --git a/media_expire.module b/media_expire.module
index efffa0d..81841dd 100644
--- a/media_expire.module
+++ b/media_expire.module
@@ -15,14 +15,13 @@ use Drupal\media\Entity\Media;
  *
  * Adds expire configuration fields to media bundle form.
  */
-function media_expire_form_media_type_form_alter(&$form, FormStateInterface $form_state, $form_id) {
+function media_expire_form_media_type_edit_form_alter(&$form, FormStateInterface $form_state, $form_id) {
 
-  /** @var \Drupal\media\MediaTypeInterface $bundle */
-  $bundle = $form_state->getStorage()['type'];
+  /** @var \Drupal\media\MediaTypeInterface $entity_type */
+  $entity_type = $form_state->getFormObject()->getEntity();
   $options = [];
   $allowed_field_types = ['datetime'];
-  $entityFieldManager = \Drupal::service('entity_field.manager');
-  foreach ($entityFieldManager->getFieldDefinitions('media', $bundle->id()) as $field_name => $field) {
+  foreach (\Drupal::service('entity_field.manager')->getFieldDefinitions('media', $entity_type->id()) as $field_name => $field) {
     if (in_array($field->getType(), $allowed_field_types) && !$field->getFieldStorageDefinition()->isBaseField()) {
       $options[$field_name] = $field->getLabel();
     }
@@ -45,17 +44,16 @@ function media_expire_form_media_type_form_alter(&$form, FormStateInterface $for
     ];
   }
   else {
-
     $form['media_expire']['enable_expiring'] = [
       '#type' => 'checkbox',
       '#title' => t('Activate media expire'),
-      '#default_value' => $bundle->getThirdPartySetting('media_expire', 'enable_expiring'),
+      '#default_value' => $entity_type->getThirdPartySetting('media_expire', 'enable_expiring'),
     ];
 
     $form['media_expire']['expire_field'] = [
       '#type' => 'select',
       '#title' => t('Expire field'),
-      '#default_value' => $bundle->getThirdPartySetting('media_expire', 'expire_field'),
+      '#default_value' => $entity_type->getThirdPartySetting('media_expire', 'expire_field'),
       '#options' => $options,
       '#description' => t('Select field that stores the expire date.'),
       '#states' => [
@@ -67,18 +65,18 @@ function media_expire_form_media_type_form_alter(&$form, FormStateInterface $for
       ],
     ];
 
-    $defaultValue = $bundle->getThirdPartySetting('media_expire', 'fallback_media');
+    $defaultValue = $entity_type->getThirdPartySetting('media_expire', 'fallback_media');
     if ($defaultValue) {
       $defaultValue = Media::load($defaultValue);
     }
 
     $form['media_expire']['fallback_media'] = [
       '#type' => 'entity_autocomplete',
-      '#title' => t('Fallback @label', ['@label' => strtolower($bundle->label())]),
+      '#title' => t('Fallback @label', ['@label' => strtolower($entity_type->label())]),
       '#default_value' => $defaultValue,
       '#target_type' => 'media',
-      '#selection_settings' => ['target_bundles' => [$bundle->id()]],
-      '#description' => t('Select a fallback @label for unpublished entities.', ['@label' => strtolower($bundle->label())]),
+      '#selection_settings' => ['target_bundles' => [$entity_type->id()]],
+      '#description' => t('Select a fallback @label for unpublished entities.', ['@label' => strtolower($entity_type->label())]),
       '#states' => [
         'visible' => [
           [