Fix for WG import

awslabs · Jun 7, 2024 · c9654fe · c9654fe
1 parent 968371b
commit c9654fe
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/python/graphstorm/wholegraph/utils.py b/python/graphstorm/wholegraph/utils.py
@@ -22,7 +22,7 @@ def is_wholegraph_embedding(data):
         is required to use wholegraph framework.
     """
     try:
-        import pylibwholegraph
+        import pylibwholegraph.torch
         return isinstance(data, pylibwholegraph.torch.WholeMemoryEmbedding)
     except ImportError:
         return False
@@ -33,7 +33,7 @@ def is_wholegraph_embedding_module(data):
     is required to use wholegraph framework.
     """
     try:
-        import pylibwholegraph
+        import pylibwholegraph.torch
         return isinstance(data, pylibwholegraph.torch.WholeMemoryEmbeddingModule)
     except:  # pylint: disable=bare-except
         return False
@@ -44,7 +44,7 @@ def is_wholegraph_optimizer(data):
     is required to use wholegraph framework.
     """
     try:
-        import pylibwholegraph
+        import pylibwholegraph.torch
         return isinstance(data, pylibwholegraph.torch.WholeMemoryOptimizer)
     except:  # pylint: disable=bare-except
         return False