From 634a76269cc01cfecea07bf6dbb6e6dd471f584d Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Sat, 27 Apr 2024 23:04:27 -0400 Subject: [PATCH] fix: `run_on_gpu` dev param --- nx_arangodb/algorithms/centrality/betweenness.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/nx_arangodb/algorithms/centrality/betweenness.py b/nx_arangodb/algorithms/centrality/betweenness.py index c8bd0f1b..9c60b148 100644 --- a/nx_arangodb/algorithms/centrality/betweenness.py +++ b/nx_arangodb/algorithms/centrality/betweenness.py @@ -1,3 +1,5 @@ +import os + from networkx.algorithms.centrality import betweenness as nx_betweenness from nx_arangodb.convert import _to_graph as _to_nx_arangodb_graph @@ -27,12 +29,14 @@ _plc="betweenness_centrality", ) def betweenness_centrality( - G, k=None, normalized=True, weight=None, endpoints=False, seed=None, run_on_gpu=True + G, k=None, normalized=True, weight=None, endpoints=False, seed=None ): print("ANTHONY: Calling betweenness_centrality from nx_arangodb") + run_on_gpu = os.environ.get("RUN_ON_GPU", "false").lower() + # 1. - if GPU_ENABLED and run_on_gpu: + if GPU_ENABLED and run_on_gpu == "true": print("ANTHONY: GPU is enabled. Using nx-cugraph bc()") if weight is not None: