diff --git a/tests/context.py b/tests/context.py
deleted file mode 100644
index 0e1b323..0000000
--- a/tests/context.py
+++ /dev/null
@@ -1,4 +0,0 @@
-import os
-import sys
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-import emate
\ No newline at end of file
diff --git a/tests/test_kpm.py b/tests/test_kpm.py
index 0e6ca7d..7923001 100644
--- a/tests/test_kpm.py
+++ b/tests/test_kpm.py
@@ -4,7 +4,8 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
-from .context import emate
+#from .context import emate
+import emate
 
 try:
     import cupy as cp
@@ -14,9 +15,12 @@
 
 
 def calc_cupykpm(W, vals, kde):
-    num_moments = 140
-    num_vecs = 140
-    extra_points = 25
+    import cupy  as cp
+    print("\nCUPY version", cp.__version__, "\n")
+
+    num_moments = 50
+    num_vecs = 50
+    extra_points = 5
     ek, rho = emate.hermitian.cupykpm(
         W.tocsr().astype("complex64"), num_moments, num_vecs, extra_points)
 
@@ -30,11 +34,15 @@ def calc_cupykpm(W, vals, kde):
     return integrate.simps(ek.get(), np.abs(rho.get()-np.exp(log_dens))) < 0.01
 
 def calc_tfkpm(W, vals, kde):
+    import tensorflow as tf
+    print("\nTF version", tf.__version__, "\n")
+
     num_moments = 40
-    num_vecs = 40
-    extra_points = 25
-    ek, rho = emate.hermitian.cupykpm(
-        W.tocsr().astype("complex64"), num_moments, num_vecs, extra_points, device="cpu:0")
+    num_vecs = 20
+    extra_points = 5
+    ek, rho = emate.hermitian.tfkpm(
+            W.tocsr().astype("complex64"), num_moments, num_vecs, extra_points,
+            device="/CPU:0")
 
     print("Saving the tfKPM plot..")
     plt.hist(vals, density=True, bins=100, alpha=.9, color="steelblue")
@@ -45,17 +53,26 @@ def calc_tfkpm(W, vals, kde):
 
     return integrate.simps(ek, np.abs(rho-np.exp(log_dens))) < 0.01
 
-def test_kpm():
+def test_tfkpm():
     n = 1000
     g = nx.erdos_renyi_graph(n , 3/n)
     W = nx.adjacency_matrix(g)
-    
+
     vals  = np.linalg.eigvals(W.todense()).real
 
     kde = KernelDensity(kernel='gaussian', bandwidth=0.1).fit(vals[:, np.newaxis])
 
     assert calc_tfkpm(W, vals, kde)
-    
+
+def test_cupykpm():
+    n = 1000
+    g = nx.erdos_renyi_graph(n , 3/n)
+    W = nx.adjacency_matrix(g)
+
+    vals  = np.linalg.eigvals(W.todense()).real
+
+    kde = KernelDensity(kernel='gaussian', bandwidth=0.1).fit(vals[:, np.newaxis])
+
     if textCupyImplementation:
         assert calc_cupykpm(W, vals, kde)
- 
\ No newline at end of file
+