Closed M-Quadra closed 4 weeks ago
import torch from torch import nn import coremltools as ct from coremltools.converters.mil.mil import types class Model(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return x.cumsum(dim=-1, dtype=torch.int32) model = Model().eval() x = torch.tensor([0.5, 1.5, 2.5]) y = model(x) print(y) traced_model = torch.jit.trace(model, (x)) mlmodel = ct.convert( traced_model, inputs=[ ct.TensorType(name="x", shape=x.shape, dtype=types.fp32), ], outputs=[ ct.TensorType(name="y"), ], ) mlmodel.save("tmp.mlpackage")
It will get the right result.
Change looks good.
CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/1509752385
Thanks for the fix @M-Quadra
It will get the right result.