Layout convolution as NHWC or NCHW only

ROCm · Dec 20, 2024 · 5be9615 · 5be9615
1 parent 52e204d
commit 5be9615
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/layout_convolution.cpp b/src/layout_convolution.cpp
@@ -40,14 +40,18 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace {
 std::vector<int64_t> get_permutation(instruction_ref ins, const layout_convolution& lc)
 {
+    std::vector<int64_t> perm(ins->get_shape().ndim());
     if(lc.channels_last)
     {
-        std::vector<int64_t> perm(ins->get_shape().ndim());
         std::iota(perm.begin() + 1, perm.end() - 1, 2);
         perm.back() = 1;
-        return perm;
     }
-    return find_permutation(ins->inputs().front()->get_shape());
+    else
+    {
+        std::iota(perm.begin(), perm.end(), 0);
+    }
+    return perm;
+
 }
 
 bool skip_layout(const shape& s)