Fix dimension check (#524)

Summary:
Fixes the assertion that `p1` and `p2` have the same last dimension. The issue was that `D` is set to equal `p2.size(2)`, and then `D` is compared to `p2.size(2)`. The fix instead compares `D` to `p1.size(2).

Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/524

Reviewed By: bottler

Differential Revision: D26008688

Pulled By: nikhilaravi

fbshipit-source-id: e32afe9da127d81b1a411d3c223b539a7400597b

pytorch3d/csrc/knn/knn.cu

@@ -505,7 +505,7 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborBackwardCuda(
   const auto D = p2.size(2);
   const auto K = idxs.size(2);
 
-  TORCH_CHECK(p2.size(2) == D, "Point sets must have the same last dimension");
+  TORCH_CHECK(p1.size(2) == D, "Point sets must have the same last dimension");
   TORCH_CHECK(idxs.size(0) == N, "KNN idxs must have the same batch dimension");
   TORCH_CHECK(
       idxs.size(1) == P1, "KNN idxs must have the same point dimension as p1");