diff --git a/src/Plugin/Derivative/EntityBlock.php b/src/Plugin/Derivative/EntityBlock.php
index da9e095..21d69f9 100644
--- a/src/Plugin/Derivative/EntityBlock.php
+++ b/src/Plugin/Derivative/EntityBlock.php
@@ -2,7 +2,6 @@
 
 namespace Drupal\entity_block\Plugin\Derivative;
 
-use Drupal\Core\Entity\ContentEntityTypeInterface;
 use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Plugin\Discovery\ContainerDeriverInterface;
 use Symfony\Component\DependencyInjection\ContainerInterface;
@@ -73,11 +72,8 @@ class EntityBlock implements ContainerDeriverInterface {
    * {@inheritdoc}
    */
   public function getDerivativeDefinitions($base_plugin_definition) {
-
     foreach ($this->entityTypeManager->getDefinitions() as $entityDefinition) {
-
-      if ($entityDefinition instanceof ContentEntityTypeInterface) {
-
+      if ($entityDefinition->hasViewBuilderClass()) {
         $delta = $entityDefinition->id();
         $this->derivatives[$delta] = [
           'category' => 'Entity Block',