[dnn/shift] optimizations for NCHW layout

2019-07-12 20:22:32 -07:00
parent 54617b4e51
commit fe42cb7142
3 changed files with 43 additions and 42 deletions
--- a/examples/cpp/shift.cpp
+++ b/examples/cpp/shift.cpp
@@ -33,7 +33,7 @@ int main() {
  triton::dnn::shift shift(B, C, 1, H, W, 1, R, S, F, 1, 1,
                           shift_h.data(), shift_w.data(),
                           numeric_t_str, numeric_t_str,
-                           op, false);
+                           op, false, triton::dnn::shift::NCHW);
  // host buffers
  std::vector<NumericT> ha(B*C*H*W);
  std::vector<NumericT> hb(C*F);