add channels

apache · Sep 5, 2019 · 9e40675 · 9e40675
1 parent d333992
commit 9e40675
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/topi/python/topi/intel_graphics/conv2d.py b/topi/python/topi/intel_graphics/conv2d.py
@@ -159,6 +159,11 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F):
 
     copy_inputs = [s for s in inputs]
     new_attrs = {k : attrs[k] for k in attrs.keys()}
+
+    if F.__name__ == 'tvm.relay.op':
+        # Derive channels for frontends (e.g ONNX) that miss "channel" field.
+        new_attrs["channels"] = inputs[1].checked_type.shape[attrs['kernel_layout'].index('O')]
+
     data, kernel = tinfo[0], tinfo[1]
     batch_size, in_channel, height, width = get_const_tuple(data.shape)