add the type of the argument.

dmlc · Jul 18, 2024 · 83099bc · 83099bc
1 parent 7fcc4b4
commit 83099bc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/python/dgl/graphbolt/impl/gpu_cached_feature.py b/python/dgl/graphbolt/impl/gpu_cached_feature.py
@@ -108,7 +108,7 @@ def wait():
 
         yield _Waiter()
 
-    def read_async_num_stages(self, ids_device):
+    def read_async_num_stages(self, ids_device: torch.device):
         """The number of stages of the read_async operation"""
         assert ids_device.type == "cuda"
         return self._fallback_feature.read_async_num_stages(ids_device) + 1

diff --git a/python/dgl/graphbolt/impl/torch_based_feature_store.py b/python/dgl/graphbolt/impl/torch_based_feature_store.py
@@ -192,7 +192,7 @@ def wait():
         else:
             yield torch.ops.graphbolt.index_select_async(self._tensor, ids)
 
-    def read_async_num_stages(self, ids_device):
+    def read_async_num_stages(self, ids_device: torch.device):
         """The number of stages of the read_async operation"""
         if ids_device.type == "cuda":
             return 1 if self.is_pinned() else 3