support for cupy

friguzzi · Dec 30, 2023 · 28526d2 · 28526d2
1 parent 1723f0a
commit 28526d2
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/prolog/liftcover.py b/prolog/liftcover.py
@@ -1,8 +1,17 @@
 
 
-import numpy as xp
+import numpy as np
 import torch
 
+try:
+    import cupy as cp
+    cp.cuda.runtime.getDeviceCount()
+    print("GPU exists!")
+    xp=cp
+except:
+    xp=np
+    print("GPU does not exist.")
+
 def ll(pr, co, zero=0.000001):
     probs=xp.array(pr)
     counts=xp.array(co)