[dnn/shift] fixed in leading dimensions for shift-conv operation

2019-07-05 17:17:22 -07:00
parent c666f71fd6
commit 3e49dbe6ab
3 changed files with 14 additions and 18 deletions
--- a/examples/python/tensorflow/shift.cpp
+++ b/examples/python/tensorflow/shift.cpp
@@ -71,6 +71,8 @@ public:
      // checks
      OP_REQUIRES(context, Ca == Cb, tensorflow::errors::InvalidArgument("operands must have the same number of channels"));
      C = Ca;
+      if(OP == triton::dnn::shift::BPROP)
+        std::swap(C, F);
    }

  }