From 47c4f51838c72a541d4916ba5cdb3281f3bc26dd Mon Sep 17 00:00:00 2001 From: Ralph Liu Date: Fri, 17 May 2024 11:20:29 -0700 Subject: [PATCH] Fix enable_batch call --- .../centrality/test_betweenness_centrality.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py b/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py index ff8859a01b1..6d1f53f7fc3 100644 --- a/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py +++ b/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py @@ -111,18 +111,13 @@ def calc_betweenness_centrality( else: edge_attr = None - G = None + G = graph_file.get_graph( + download=True, + create_using=cugraph.Graph(directed=directed), + ignore_weights=not edgevals, + ) if multi_gpu_batch: - G = graph_file.get_dask_graph( - create_using=cugraph.Graph(directed=directed), ignore_weights=not edgevals - ) G.enable_batch() - else: - G = graph_file.get_graph( - download=True, - create_using=cugraph.Graph(directed=directed), - ignore_weights=not edgevals, - ) M = G.to_pandas_edgelist().rename( columns={"src": "0", "dst": "1", "wgt": edge_attr}