fix: run_on_gpu dev param

arangodb · Apr 28, 2024 · 634a762 · 634a762
1 parent 4797eb9
commit 634a762
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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: