diff --git a/src/EntityHelperTrait.php b/src/EntityHelperTrait.php
index 08030ce..908502a 100644
--- a/src/EntityHelperTrait.php
+++ b/src/EntityHelperTrait.php
@@ -213,8 +213,8 @@ trait EntityHelperTrait {
       $build['#attributes']['data-caption'] = $context['data-caption'];
     }
 
-    // If this is an image, the image_formatter expects #item_attributes.
-    if ($build['#theme'] == 'image_formatter') {
+    // If this is an image, the image_formatter template expects #item_attributes.
+    if (!empty($build['#theme']) && !empty($build['#attributes']) && $build['#theme'] == 'image_formatter') {
       $build['#item_attributes'] =  $build['#attributes'];
     }