diff --git a/pytorch3d/csrc/nearest_neighbor_points/nearest_neighbors_points_cpu.cpp b/pytorch3d/csrc/nearest_neighbor_points/nearest_neighbors_points_cpu.cpp index 0fa80488..3dd373b9 100644 --- a/pytorch3d/csrc/nearest_neighbor_points/nearest_neighbors_points_cpu.cpp +++ b/pytorch3d/csrc/nearest_neighbor_points/nearest_neighbors_points_cpu.cpp @@ -3,36 +3,36 @@ #include at::Tensor NearestNeighborIdxCpu(at::Tensor p1, at::Tensor p2) { - const int N = p1.size(0); - const int P1 = p1.size(1); - const int D = p1.size(2); - const int P2 = p2.size(1); + const int N = p1.size(0); + const int P1 = p1.size(1); + const int D = p1.size(2); + const int P2 = p2.size(1); - auto long_opts = p1.options().dtype(torch::kInt64); - torch::Tensor out = torch::empty({N, P1}, long_opts); + auto long_opts = p1.options().dtype(torch::kInt64); + torch::Tensor out = torch::empty({N, P1}, long_opts); - auto p1_a = p1.accessor(); - auto p2_a = p2.accessor(); - auto out_a = out.accessor(); + auto p1_a = p1.accessor(); + auto p2_a = p2.accessor(); + auto out_a = out.accessor(); - for (int n = 0; n < N; ++n) { - for (int i1 = 0; i1 < P1; ++i1) { - // TODO: support other floating-point types? - float min_dist = -1; - int64_t min_idx = -1; - for (int i2 = 0; i2 < P2; ++i2) { - float dist = 0; - for (int d = 0; d < D; ++d) { - float diff = p1_a[n][i1][d] - p2_a[n][i2][d]; - dist += diff * diff; - } - if (min_dist == -1 || dist < min_dist) { - min_dist = dist; - min_idx = i2; - } - } - out_a[n][i1] = min_idx; + for (int n = 0; n < N; ++n) { + for (int i1 = 0; i1 < P1; ++i1) { + // TODO: support other floating-point types? + float min_dist = -1; + int64_t min_idx = -1; + for (int i2 = 0; i2 < P2; ++i2) { + float dist = 0; + for (int d = 0; d < D; ++d) { + float diff = p1_a[n][i1][d] - p2_a[n][i2][d]; + dist += diff * diff; } + if (min_dist == -1 || dist < min_dist) { + min_dist = dist; + min_idx = i2; + } + } + out_a[n][i1] = min_idx; } - return out; + } + return out; } diff --git a/scripts/parse_tutorials.py b/scripts/parse_tutorials.py index 550c3975..e4ec568e 100644 --- a/scripts/parse_tutorials.py +++ b/scripts/parse_tutorials.py @@ -9,7 +9,6 @@ import nbformat from bs4 import BeautifulSoup from nbconvert import HTMLExporter, ScriptExporter - TEMPLATE = """const CWD = process.cwd(); const React = require('react'); @@ -43,7 +42,9 @@ def gen_tutorials(repo_dir: str) -> None: Also create ipynb and py versions of tutorial in Docusaurus site for download. """ - with open(os.path.join(repo_dir, "website", "tutorials.json"), "r") as infile: + with open( + os.path.join(repo_dir, "website", "tutorials.json"), "r" + ) as infile: tutorial_config = json.loads(infile.read()) tutorial_ids = {x["id"] for v in tutorial_config.values() for x in v} @@ -52,7 +53,9 @@ def gen_tutorials(repo_dir: str) -> None: print("Generating {} tutorial".format(tid)) # convert notebook to HTML - ipynb_in_path = os.path.join(repo_dir, "docs", "tutorials", "{}.ipynb".format(tid)) + ipynb_in_path = os.path.join( + repo_dir, "docs", "tutorials", "{}.ipynb".format(tid) + ) with open(ipynb_in_path, "r") as infile: nb_str = infile.read() nb = nbformat.reads(nb_str, nbformat.NO_CONVERT) @@ -105,7 +108,10 @@ if __name__ == "__main__": description="Generate JS, HTML, ipynb, and py files for tutorials." ) parser.add_argument( - "--repo_dir", metavar="path", required=True, help="Pytorch3D repo directory." + "--repo_dir", + metavar="path", + required=True, + help="Pytorch3D repo directory.", ) args = parser.parse_args() - gen_tutorials(args.repo_dir) \ No newline at end of file + gen_tutorials(args.repo_dir) diff --git a/tests/test_nearest_neighbor_points.py b/tests/test_nearest_neighbor_points.py index 964b5a9a..43bcafec 100644 --- a/tests/test_nearest_neighbor_points.py +++ b/tests/test_nearest_neighbor_points.py @@ -43,21 +43,21 @@ class TestNearestNeighborPoints(unittest.TestCase): """ Test cuda output vs naive python implementation. """ - device = torch.device('cuda:0') + device = torch.device("cuda:0") self._test_nn_helper(device) def test_nn_cpu(self): """ Test cpu output vs naive python implementation """ - device = torch.device('cpu') + device = torch.device("cpu") self._test_nn_helper(device) @staticmethod def bm_nn_points_cpu_with_init( N: int = 4, D: int = 4, P1: int = 128, P2: int = 128 ): - device = torch.device('cpu') + device = torch.device("cpu") x = torch.randn(N, P1, D, device=device) y = torch.randn(N, P2, D, device=device)