diff --git a/modules/conditional_fields/src/Plugin/conditional_fields/handler/EntityReference.php b/modules/conditional_fields/src/Plugin/conditional_fields/handler/EntityReference.php
index 229f0dd..8d65450 100644
--- a/modules/conditional_fields/src/Plugin/conditional_fields/handler/EntityReference.php
+++ b/modules/conditional_fields/src/Plugin/conditional_fields/handler/EntityReference.php
@@ -3,7 +3,7 @@
 namespace Drupal\conditional_fields\Plugin\conditional_fields\handler;
 
 use Drupal\conditional_fields\ConditionalFieldsHandlerBase;
-use Drupal\node\Entity\Node;
+use Drupal\Core\Entity;
 
 /**
  * Provides states handler for entity reference fields.
@@ -24,14 +24,16 @@ public function statesHandler($field, $field_info, $options) {
     switch ($values_set) {
       case CONDITIONAL_FIELDS_DEPENDENCY_VALUES_WIDGET:
         $target_id = $this->getWidgetValue($options['value_form']);
+        $entity_type = $field['#target_type'];
         if (!empty($target_id)) {
-          $node = Node::load($target_id);
+          //TODO: this should be injected
+          $entity = \Drupal::entityTypeManager()->getStorage($entity_type)->load($target_id);
           // Create an array of valid formats of title for autocomplete.
           $state[$options['state']][$options['selector']] = [
-            // Node title (nid).
-            ['value' => $node->label() . ' (' . $node->id() . ')'],
-            // Node title.
-            ['value' => $node->label()]
+            // Entity title (nid).
+            ['value' => $entity->label() . ' (' . $entity->id() . ')'],
+            // Entity title.
+            ['value' => $entity->label()]
           ];
         }
         break;