fix tensorflow relay converter for conv3d_transpose op

apache · Mar 23, 2023 · 3b1c52a · 3b1c52a
1 parent c16a59d
commit 3b1c52a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/python/tvm/relay/frontend/tensorflow_ops.py b/python/tvm/relay/frontend/tensorflow_ops.py
@@ -693,7 +693,10 @@ def _impl(inputs, attr, params, mod):
             raise tvm.error.OpAttributeInvalid(msg.format(attr["padding"]))
 
         if "kernel_layout" not in attr:
-            attr["kernel_layout"] = "DHWIO" if attr["data_format"] == "NDHWC" else "OIDHW"
+            if opname == "conv":
+                attr["kernel_layout"] = "DHWIO" if attr["data_format"] == "NDHWC" else "OIDHW"
+            elif opname == "conv_transpose":
+                attr["kernel_layout"] = "DHWOI" if attr["data_format"] == "NDHWC" else "IODHW"
 
         use_bias = len(inputs) == (3 if opname != "conv_transpose" else 4)
         channel_axis = 1 if attr["data_format"] == "NCDHW" else 4