linting

dmlc · Jul 18, 2024 · 10ba72c · 10ba72c
1 parent 1d49dc8
commit 10ba72c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/python/dgl/graphbolt/impl/gpu_cached_feature.py b/python/dgl/graphbolt/impl/gpu_cached_feature.py
@@ -2,11 +2,9 @@
 
 import torch
 
-from ..base import get_device_to_host_uva_stream, get_host_to_device_uva_stream
 from ..feature_store import Feature
 
 from .gpu_cache import GPUCache
-from .torch_based_feature_store import DiskBasedFeature, TorchBasedFeature
 
 __all__ = ["GPUCachedFeature"]
 
@@ -104,6 +102,7 @@ def read_async(self, ids: torch.Tensor):
         class _Waiter:
             @staticmethod
             def wait():
+                """Returns the stored value when invoked."""
                 return values
 
         yield _Waiter()

diff --git a/python/dgl/graphbolt/impl/torch_based_feature_store.py b/python/dgl/graphbolt/impl/torch_based_feature_store.py
@@ -154,6 +154,7 @@ def read_async(self, ids: torch.Tensor):
             class _Waiter:
                 @staticmethod
                 def wait():
+                    """Returns the stored value when invoked."""
                     values_copy_event.wait()
                     return values
 
@@ -185,6 +186,7 @@ def wait():
             class _Waiter:
                 @staticmethod
                 def wait():
+                    """Returns the stored value when invoked."""
                     values_copy_event.wait()
                     return values_cuda
 
@@ -423,6 +425,7 @@ def read_async(self, ids: torch.Tensor = None):
             class _Waiter:
                 @staticmethod
                 def wait():
+                    """Returns the stored value when invoked."""
                     values_copy_event.wait()
                     return values_cuda