diff --git a/src/PluginTypeManager.php b/src/PluginTypeManager.php
index 5a5ef314774ecd095d736581d44f60160d843394..0f6f8912405775fea2284af44a253ff73e9e20e3 100644
--- a/src/PluginTypeManager.php
+++ b/src/PluginTypeManager.php
@@ -6,7 +6,6 @@ use Drupal\Core\Config\Entity\ConfigEntityBase;
 use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\hierarchy_manager\Plugin\HmDisplayPluginManager;
 use Drupal\hierarchy_manager\Plugin\HmSetupPluginManager;
-use DrupalCodeGenerator\Command\Drupal_8\Module\ConfigurationEntity;
 
 class PluginTypeManager {
   
@@ -67,7 +66,7 @@ class PluginTypeManager {
   /**
    * Get a display plugin instance according to a setup plugin.
    * 
-   * @param \Drupal\Core\Config\Entity\ConfigEntityBase; $display_profile
+   * @param \Drupal\Core\Config\Entity\ConfigEntityBase $display_profile
    *   Display profile entity.
    * @return NULL|object
    *   The display plugin instance.
@@ -87,7 +86,7 @@ class PluginTypeManager {
    *
    * @param string $setup_plugin_id
    *   setup plugin ID.
-   * @return NULL|\DrupalCodeGenerator\Command\Drupal_8\Module\ConfigurationEntity
+   * @return NULL|\Drupal\Core\Config\Entity\ConfigEntityBase
    *   The display profile entity.
    */
   public function getDisplayProfile(string $setup_plugin_id) {