Add torch one_hot op - Githubissues

Example:

import torch
from torch import nn
from torch.nn import functional as F
from typing import Final
from cp import coremltools as ct
from cp.coremltools.converters.mil.mil import types

class Model(nn.Module):
  def forward(self, x: torch.LongTensor) -> torch.Tensor:
      return F.one_hot(x, num_classes=10)

model = Model().eval()
x = torch.arange(10)
traced_model = torch.jit.trace(model, (x))

var_dim: Final[ct.RangeDim] = ct.RangeDim(1, 1_000)
mlmodel = ct.convert(
    traced_model,
    inputs=[
        ct.TensorType(name="x", shape=ct.Shape([var_dim]), dtype=types.int32),
    ],
    outputs=[
        ct.TensorType(name="y")
    ],
)
mlmodel.save("tmp.mlpackage")

apple / coremltools

Add torch one_hot op #2358