diff --git a/gsaml.module b/gsaml.module
index a149f20..b3a7660 100644
--- a/gsaml.module
+++ b/gsaml.module
@@ -1,5 +1,6 @@
 <?php
 
+use Drupal\user\UserInterface;
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\gsaml\GSAML;
 use Drupal\Core\Entity\EntityFormInterface;
@@ -57,7 +58,7 @@ function _gsaml_entity_update($form, FormStateInterface $form_state) {
 /**
  * Associate user to group/roles on login.
  */
-function gsaml_simplesamlphp_auth_user_attributes(\Drupal\user\UserInterface $account, $attributes) {
+function gsaml_simplesamlphp_auth_user_attributes(UserInterface $account, $attributes) {
   $config_factory = \Drupal::configFactory();
   $config = $config_factory->getEditable('gsaml.settings');
   $field_name = $config->get('user_field');
diff --git a/src/GSAML.php b/src/GSAML.php
index b43f78e..8ea6d2b 100644
--- a/src/GSAML.php
+++ b/src/GSAML.php
@@ -509,7 +509,7 @@ class GSAML {
       $nid = self::getNextEntityWithoutGroupLargerThan($context['sandbox']['curr_node'], $type);
       $context['sandbox']['curr_node'] = $nid;
       if ($type == 'media') {
-        $entity = \Drupal\media\Entity\Media::load($nid);
+        $entity = Media::load($nid);
       }
       else {
         $entity = Node::load($nid);