diff --git a/merlin/models/tf/distributed/backend.py b/merlin/models/tf/distributed/backend.py
index 4e754106ed..91f9d7707c 100644
--- a/merlin/models/tf/distributed/backend.py
+++ b/merlin/models/tf/distributed/backend.py
@@ -22,7 +22,7 @@
 
 if HAS_GPU:
     try:
-        from sparse_operation_kit import sparse_operation_kit as sok  # noqa: F401
+        import sparse_operation_kit as sok  # noqa: F401
 
         sok_installed = True
     except (ImportError, tf.errors.NotFoundError):