Fix tests

Author: nicolas-chaulet

cpu/include/neighbors.h

@@ -20,4 +20,4 @@ int batch_nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppo
 
 template <typename scalar_t>
 void nanoflann_knn_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
-                        vector<long>& neighbors_indices, vector<float>& dists, int k);
+                             vector<long>& neighbors_indices, vector<float>& dists, int k);

cpu/src/knn.cpp

@@ -16,7 +16,7 @@ std::pair<at::Tensor, at::Tensor> _single_batch_knn(at::Tensor support, at::Tens
     std::vector<float> neighbors_dists(query.size(0) * k, -1);
 
     auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
-    auto options_dist = torch::TensorOptions().dtype(query.scalar_type()).device(torch::kCPU);
+    auto options_dist = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
     AT_DISPATCH_ALL_TYPES(query.scalar_type(), "knn", [&] {
         auto data_q = query.DATA_PTR<scalar_t>();
         auto data_s = support.DATA_PTR<scalar_t>();

test/test_interpolate.py

@@ -36,7 +36,7 @@ def test_grad(self):
         dist_recip = 1.0 / (dist + 1e-8)
         norm = torch.sum(dist_recip, dim=2, keepdim=True)
         weight = dist_recip / norm
-        input = (x, idx, weight)
+        input = (x, idx, weight.double())
         test = gradcheck(three_interpolate, input, eps=1e-6, atol=1e-4)
 
 

test/test_knn.py

@@ -12,8 +12,8 @@
 
 class TestKnn(unittest.TestCase):
     def test_cpu(self):
-        support = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]]])
-        query = torch.tensor([[[0, 0, 0]]])
+        support = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]]]).float()
+        query = torch.tensor([[[0, 0, 0]]]).float()
 
         idx, dist = knn(support, query, 3)
         torch.testing.assert_allclose(idx, torch.tensor([[[0, 1, 2]]]))