diff --git a/includes/entity_embed.display.inc b/includes/entity_embed.display.inc
index f250582..6bb8ab8 100644
--- a/includes/entity_embed.display.inc
+++ b/includes/entity_embed.display.inc
@@ -55,12 +55,38 @@ function entity_embed_render_entity($entity_type, $entity, array $context = arra
     $context
   );
 
-  // Maintain data-align if it is there.
+  // Maintain data-align if it is there
+  // Idea from https://www.drupal.org/node/2676804
+  // Get list of entities from cache if possible; if not, use entity_get_info() and cache the result
+  if (!($entities = cache_get('entity_embed_types'))) {
+    $entities = array_keys(entity_get_info());
+    cache_set('entity_embed_types', $entities);
+  } else {
+    $entities = $entities->data;
+  }
   if (isset($context['data-align'])) {
-    $build['#attributes']['data-align'] = $context['data-align'];
+    foreach($entities as $entity_type) {
+      if (isset($build[$entity_type])) {
+        foreach ($build[$entity_type] as $entity_id => $entity_array) {
+          if (is_array($entity_array)) {
+            $build[$entity_type][$entity_id]['#attributes']['data-align'] = $context['data-align'];
+            $build[$entity_type][$entity_id]['#type'] = 'container';
+          }
+        }
+      }
+    }
   }
   elseif ((isset($context['class']))) {
-    $build['#attributes']['class'][] = $context['class'];
+    foreach($entities as $entity_type) {
+      if (isset($build[$entity_type])) {
+        foreach ($build[$entity_type] as $entity_id => $entity_array) {
+          if (is_array($entity_array)) {
+            $build[$entity_type][$entity_id]['#attributes']['class'][] = $context['class'];
+            $build[$entity_type][$entity_id]['#type'] = 'container';
+          }
+        }
+      }
+    }
   }
 
   // Maintain data-caption if it is there.