Add support for conv_transpose and add aliases

apple · Nov 10, 2023 · 906f052 · 906f052
1 parent b2f7190
commit 906f052
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py
@@ -930,7 +930,7 @@ def linear(context, node):
     context.add(res, torch_name=node.name)
 
 
-@register_torch_op(torch_alias=["conv2d", "convolution"])
+@register_torch_op(torch_alias=["convolution", "conv1d", "conv2d", "conv3d", "conv_transpose1d", "conv_transpose2d", "conv_transpose3d"])
 def _convolution(context, node):
     inputs = _get_inputs(context, node)
 
@@ -968,6 +968,11 @@ def _convolution(context, node):
         transposed = inputs[6].val
         out_pad = inputs[7].val
         group = inputs[8]
+    elif len(inputs) == 8:
+        transposed = True
+        out_pad = inputs[5].val
+        dilations = inputs[7]
+        group = inputs[6]
     elif len(inputs) == 7:
         transposed = False
         group = inputs[6]