diff --git a/amp.module b/amp.module
index 1b8e922..e27db3a 100644
--- a/amp.module
+++ b/amp.module
@@ -37,6 +37,20 @@ function amp_help($route_name, RouteMatchInterface $route_match) {
 }
 
 /**
+ * Implements hook_entity_view_mode_alter().
+ */
+function amp_entity_view_mode_alter(&$view_mode, EntityInterface $entity, $context) {
+
+  // Get the AMP Context service.
+  $amp_context = \Drupal::service('router.amp_context');
+
+  // If on AMP route, and in full view mode, switch to AMP view mode.
+  if ($amp_context->isAmpRoute() && $view_mode == 'full') {
+    $view_mode = 'amp';
+  }
+}
+
+/**
  * Implements hook_entity_view_alter().
  */
 function amp_entity_view_alter(array &$build, EntityInterface $entity, EntityViewDisplayInterface $display) {
@@ -50,10 +64,7 @@ function amp_entity_view_alter(array &$build, EntityInterface $entity, EntityVie
   // Check if entity is a node with an enabled amp view mode.
   if ($entity->getEntityType()->id() == 'node' && isset($view_modes['amp'])) {
 
-    // If on AMP route, and in full view mode, switch to AMP view mode.
-    if ($amp_context->isAmpRoute() && $build['#view_mode'] == 'full') {
-      $build['#view_mode'] = 'amp';
-
+    if ($build['#view_mode'] == 'amp') {
       // Otherwise adding 'amp' query parameter in URL will have no effect.
       $build['#cache']['contexts'] = Cache::mergeContexts($build['#cache']['contexts'], ['url.query_args:amp']);
       // Otherwise adding 'warnfix' query parameter in URL will have no effect.
@@ -65,6 +76,7 @@ function amp_entity_view_alter(array &$build, EntityInterface $entity, EntityVie
     if (!empty($build['#attached']['html_head_link'])) {
       foreach ($build['#attached']['html_head_link'] as $key => $config) {
         if ($build['#view_mode'] == 'amp') {
+
           if ($config[0]['rel'] != 'canonical' && $config[0]['rel'] != 'shortlink') {
             unset($build['#attached']['html_head_link'][$key]);
           }