Merge pull request #3282 from rcurtin/max-pooling-identical-elements

Handle MaxPooling index computation for identical elements
2022-10-09 19:03:15 -04:00
parent 2aec3ce8e9 a4dab0e427
commit 1aca420224
1 changed files with 2 additions and 3 deletions
@@ -42,9 +42,8 @@ class MaxPoolingRule
  std::tuple<size_t, typename MatType::elem_type> PoolingWithIndex(
      const MatType& input)
  {
-    const typename MatType::elem_type maxVal =
-        arma::max(arma::vectorise(input));
-    const size_t index = arma::as_scalar(arma::find(input == maxVal, 1));
+    const size_t index = input.index_max();
+    const typename MatType::elem_type maxVal = input[index];

    return std::tuple<size_t, typename MatType::elem_type>(index, maxVal);
  }