make the code easier to read.

dmlc · Aug 11, 2024 · be4009d · be4009d
1 parent 6e2970e
commit be4009d
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/tests/python/pytorch/graphbolt/test_dataloader.py b/tests/python/pytorch/graphbolt/test_dataloader.py
@@ -104,13 +104,13 @@ def test_gpu_sampling_DataLoader(
     for i in range(2):
         datapipe = dgl.graphbolt.ItemSampler(itemset, batch_size=B)
         datapipe = datapipe.copy_to(F.ctx())
-        kwargs = {}
-        if i == 0:
-            kwargs = {
-                "overlap_fetch": overlap_graph_fetch,
-                "num_gpu_cached_edges": num_gpu_cached_edges,
-                "gpu_cache_threshold": gpu_cache_threshold,
-            }
+        kwargs = {
+            "overlap_fetch": overlap_graph_fetch,
+            "num_gpu_cached_edges": num_gpu_cached_edges,
+            "gpu_cache_threshold": gpu_cache_threshold,
+        }
+        if i != 0:
+            kwargs = {}
         datapipe = getattr(dgl.graphbolt, sampler_name)(
             datapipe,
             graph,