add test to cover GPU (facebookresearch#4130)

Summary:
Pull Request resolved: facebookresearch#4130

same as title

Reviewed By: asadoughi

Differential Revision: D68388863

fbshipit-source-id: 4ebb38d8454bf95733c918950d7d8d3b22e00d5d

Author: junjieqi

tests/test_contrib.py

@@ -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)
 
@@ -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
@@ -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_PRODUCT, ngpu=-1)
+
 
 class TestDatasets(unittest.TestCase):
     """here we test only the synthetic dataset. Datasets that require