Skip to content

Commit

Permalink
add test to cover GPU
Browse files Browse the repository at this point in the history
Summary: same as title

Differential Revision: D68388863
  • Loading branch information
junjieqi authored and facebook-github-bot committed Jan 19, 2025
1 parent 2744f22 commit 72387e1
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions tests/test_contrib.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@

class TestComputeGT(unittest.TestCase):

def do_test_compute_GT(self, metric=faiss.METRIC_L2):
def do_test_compute_GT(self, metric=faiss.METRIC_L2, ngpu=0):
d = 64
xt, xb, xq = get_dataset_2(d, 0, 10000, 100)

Expand All @@ -50,7 +50,7 @@ def matrix_iterator(xb, bs):
yield xb[i0:i0 + bs]

Dnew, Inew = knn_ground_truth(
xq, matrix_iterator(xb, 1000), 10, metric, ngpu=0)
xq, matrix_iterator(xb, 1000), 10, metric, ngpu=ngpu)

np.testing.assert_array_equal(Iref, Inew)
# decimal = 4 required when run on GPU
Expand All @@ -62,6 +62,12 @@ def test_compute_GT(self):
def test_compute_GT_ip(self):
self.do_test_compute_GT(faiss.METRIC_INNER_PRODUCT)

def test_compute_GT_gpu(self):
self.do_test_compute_GT(ngpu=-1)

def test_compute_GT_ip_gpu(self):
self.do_test_compute_GT(faiss.METRIC_INNER_PRODUCTm, ngpu=-1)


class TestDatasets(unittest.TestCase):
"""here we test only the synthetic dataset. Datasets that require
Expand Down

0 comments on commit 72387e1

Please sign in to comment.