diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 33d2e6105..939eae38b 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -7277,6 +7277,25 @@ def forward(self, spans, embedding): compute_unit=compute_unit, ) + def test_cast_length_1_constant_tensor(self): + class TestModel(torch.nn.Module): + def forward(self, x): + b, c, h, w = x.shape + return x.reshape(b, c, int(h * w)) + + x = torch.rand(1, 2, 4, 4) + torch_model = export_torch_model_to_frontend( + TestModel().eval(), + (x,), + TorchFrontend.TORCHSCRIPT, + ) + ct.convert( + torch_model, + inputs=[ct.TensorType(shape=x.shape)], + convert_to="neuralnetwork", + minimum_deployment_target=ct.target.iOS14, + ) + @pytest.mark.parametrize( "compute_unit, backend, frontend", itertools.product(compute_units, backends, frontends),