Coremltools inconsistency between torch.jit.script and torch.jit.trace conversion

[JRGit4UE]

🐞Describe the bug

Converting a simple test model in Pytorch with only a Conv2d() operation in its forward() method, the conversion succeeds with torch.jit.trace but fails with torch.jit.script, which seems to be not consistent.

Trace

If applicable, please paste the error trace. Trying to convert the following example fails with an assertion

assert str(node.output().type()) == "Tensor"

in coremltools\converters\mil\frontend\torch\converter.py function '_check_is_tensor()' for parameter bias in Conv2d() which is defined as Optional (c10::optional<Tensor>& bias_opt) in
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp

at::Tensor conv2d(
    const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups)

When compiling the model as script
function _check_is_tensor() runs into an assertion, as Optional[Tensor] != Tensor
This error does not occur when compiling the model via trace, but that is not what I need.

To Reproduce

• If a python script can reproduce the error, please paste the code snippet

class Test(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1,1,3, bias=True)
    def forward(self, x:Tensor)->Tensor:
        y = self.conv(x)
        return y

import coremltools as ct
with torch.no_grad():
    torch.set_default_tensor_type('torch.FloatTensor')
    b = torch.rand(1,1,4,4)
    t = Test()
    t.eval()
    sd = t.state_dict()
    sd['conv.bias'] = torch.zeros(1)
    t.load_state_dict(sd)
    res = t(b)

script = torch.jit.script(t) # does not work
mlmodel = ct.convert(
    script,
    source="pytorch",
    inputs=[ct.TensorType(name='x', shape=b.shape)],
    convert_to="mlprogram",
    debug=False
)

    script = torch.jit.trace(t,b) # does work, but I would need  torch.jit.script
    mlmodel = ct.convert(
        script,
        source="pytorch",
        inputs=[ct.TensorType(name='x', shape=b.shape)],
        convert_to="mlpackage",
        debug=False
    )

[TobyRoseman]

@JRGit4UE - thanks for reporting the issue. I've made some minor changes to your code, in order to reproduce the issue.

This is related to #1367

[oulcan]

I have the same problem also, are there any workaround/solution for Optional[Tensor] type for coreml conversion?

[JRGit4UE]

In version 5.2.0 the issue still exists. May I suggest to adjust the coremltools documentation to be make clear - it is not production ready code - as the documentation implies -

but "experimental" when it comes to convert Pytorch scripted models

, so that Pytorch people that want to use coremltools the scripted way will know: it's at their own risk and may cost them a lot of time.