Merge branch 'branch-24.08' of https://github.com/rapidsai/cugraph in…

…to fea_biased_sampling
rapidsai · Jul 1, 2024 · 8581190 · 8581190
2 parents c2e9b5f + 70c236c
commit 8581190
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py b/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py
@@ -15,6 +15,7 @@
 import argparse
 import tempfile
 import os
+import warnings
 
 from typing import Optional, Tuple, Dict
 
@@ -159,6 +160,10 @@ def parse_args():
         args.dataset, args.dataset_root
     )
 
+    if "CI_RUN" in os.environ and os.environ["CI_RUN"] == "1":
+        warnings.warn("Pruning test dataset for CI run.")
+        split_idx["test"] = split_idx["test"][:1000]
+
     with tempfile.TemporaryDirectory(dir=args.tempdir_root) as samples_dir:
         loader_kwargs = {
             "data": data,