Fix channel check for Conv2dTransposed

apache · Nov 6, 2021 · b34f268 · b34f268
1 parent 4fdd06e
commit b34f268
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/relay/op/nn/convolution.h b/src/relay/op/nn/convolution.h
@@ -1103,7 +1103,7 @@ bool Conv2DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& a
           << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
     }
     if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(param->channels, wshape[1]))
+      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
           << "Conv2DTransposed: shape of weight is inconsistent with channels, "
           << " channels=" << param->channels << " wshape=" << Array<IndexExpr>(wshape);
     }