make fps faster on CPU

4b01cc80 · rusty1s · 3d96a5d6 · 4b01cc80
Commit 4b01cc80 authored Dec 03, 2020 by rusty1s
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

csrc/cpu/fps_cpu.cpp csrc/cpu/fps_cpu.cpp +1 -1

No files found.
--- a/csrc/cpu/fps_cpu.cpp
+++ b/csrc/cpu/fps_cpu.cpp
@@ -3,7 +3,7 @@
 #include "utils.h"

 inline torch::Tensor get_dist(torch::Tensor x, int64_t idx) {
-  return (x - x[idx]).norm(2, 1);
+  return (x - x[idx]).pow_(2).sum(1);
 }

 torch::Tensor fps_cpu(torch::Tensor src, torch::Tensor ptr, double ratio,