Merge pull request #6 from alexbarghi-nv/add-val-limit

[IMP] Limit the Test Data Size when Running CI in gcn_dist_sg.py

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,