image_colorizer_full_version_convert.py

from colorizers import *
import coremltools as ct
import torch

siggraph17_model = siggraph17(pretrained=True).eval() # ➊

example_input = torch.rand(1, 1, 256, 256)
traced_model = torch.jit.trace(siggraph17_model, example_input) # ➋

coreml_model = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input1", shape=(1, 1, 256, 256))]
)
coreml_model.save("coremlColorizer.mlmodel") # ➌