pytorch model is transformed into coreml model expect a tensor of type integer, but actually receive a tensor of type floating point.

[yuqilol]

When I was converting pytorch model into coreml model, an error occurred, indicating that the coreml model required int32 parameters but input fp32 parameters. I checked the previous error report and tried to modify it, but it seemed to be different from my error feeling. How should I improve my code to match the result?
Prediction Mismatch between Pytorch and CoreML #1428
Wrong output from PyTorch converted model #1486
Model converts fine to neuralnetwork but produces all nan values for mlprogram #1809

python: 3.10.14
pytorch: 2.2.2
numpy: 1.26.4
coremltools: 7.1
Ubuntu 18.04.6

my code

import torch
import time
import torch
import coremltools as ct
from torchvision import transforms
from models.yowo.yowo import YOWO
from config.yowo_v2_config import yowo_v2_config
from coremltools.converters.mil.mil import types
import numpy as np
yowo_v2_config["yowo_v2_tiny"]
path_to_ckpt = '/data1/home/pengyuling/yowov2/weights/yowo_v2_tiny_epoch_986.pth'
device = "cpu"
transform = transforms.Compose([
transforms.Resize([224,224]),
transforms.ToTensor(),
])
model = YOWO(
cfg = yowo_v2_config["yowo_v2_tiny"],
device = device,
num_classes = 8,
conf_thresh = 0.,
nms_thresh = 0.2,
topk = 3,
trainable = False,
multi_hot = True )
device = torch.device(device)
checkpoint = torch.load(path_to_ckpt, map_location='cpu')
checkpoint_state_dict = checkpoint.pop("model")
model_state_dict = model.state_dict()
for k in list(checkpoint_state_dict.keys()):
if k in model_state_dict:
shape_model = tuple(model_state_dict[k].shape)
shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
if shape_model != shape_checkpoint:
checkpoint_state_dict.pop(k)
else:
checkpoint_state_dict.pop(k)
print(k)
model.load_state_dict(checkpoint_state_dict)
model.eval()
x = torch.rand(1, 3, 16, 224, 224)
traced_model = torch.jit.trace(model, x)
scripted_model = torch.jit.script(traced_model)
mlmodel = ct.convert(
scripted_model, source="pytorch",
convert_to="mlprogram",
inputs=[ct.TensorType(shape=x.shape)]
)
mlmodel.save('yowo_v2_tiny_epoch_986.mlpackage')

my error

Tuple detected at graph output. This will be flattened in the converted model.
Converting PyTorch Frontend ==> MIL Ops: 66%|███████████████████████████████████▍ | 1717/2613 [00:03<00:01, 473.30 ops/s]Saving value type of int64 into a builtin type of int32, might lose precision!
Saving value type of int64 into a builtin type of int32, might lose precision!
Converting PyTorch Frontend ==> MIL Ops: 76%|█████████████████████████████████████████▏ | 1992/2613 [00:03<00:01, 592.53 ops/s]Saving value type of int64 into a builtin type of int32, might lose precision!
Saving value type of int64 into a builtin type of int32, might lose precision!
Converting PyTorch Frontend ==> MIL Ops: 87%|███████████████████████████████████████████████▏ | 2281/2613 [00:04<00:00, 585.92 ops/s]Saving value type of int64 into a builtin type of int32, might lose precision!
Saving value type of int64 into a builtin type of int32, might lose precision!
Converting PyTorch Frontend ==> MIL Ops: 93%|██████████████████████████████████████████████████▍ | 2439/2613 [00:04<00:00, 582.65 ops/s]
Traceback (most recent call last):
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/runpy.py", line 196, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/runpy.py", line 86, in _run_code
exec(code, run_globals)
File "/data1/home/pengyuling/.vscode-server/extensions/ms-python.debugpy-2024.0.0/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/main.py", line 39, in
cli.main()
File "/data1/home/pengyuling/.vscode-server/extensions/ms-python.debugpy-2024.0.0/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
run()
File "/data1/home/pengyuling/.vscode-server/extensions/ms-python.debugpy-2024.0.0/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
runpy.run_path(target, run_name="main")
File "/data1/home/pengyuling/.vscode-server/extensions/ms-python.debugpy-2024.0.0/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
return _run_module_code(code, init_globals, run_name,
File "/data1/home/pengyuling/.vscode-server/extensions/ms-python.debugpy-2024.0.0/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
_run_code(code, mod_globals, init_globals,
File "/data1/home/pengyuling/.vscode-server/extensions/ms-python.debugpy-2024.0.0/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
exec(code, run_globals)
File "/data1/home/pengyuling/yowov2/export_coreml.py", line 53, in
mlmodel = ct.convert(
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/_converters_entry.py", line 574, in convert
mlmodel = mil_convert(
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 286, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 108, in call
return load(*args, **kwargs)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 80, in load
return _perform_torch_convert(converter, debug)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 99, in _perform_torch_convert
prog = converter.convert()
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 519, in convert
convert_nodes(self.context, self.graph)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 88, in convert_nodes
add_op(context, node)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 3789, in index
x = mb.gather(x=x, indices=indices, axis=axis, name=node.name)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/mil/ops/registry.py", line 182, in add_op
return cls._add_op(op_cls_to_add, **kwargs)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/mil/builder.py", line 168, in _add_op
new_op = op_cls(**kwargs)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/mil/operation.py", line 190, in init
self._validate_and_set_inputs(input_kv)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/mil/operation.py", line 503, in _validate_and_set_inputs
self.input_spec.validate_inputs(self.name, self.op_type, input_kvs)
File "/data1/home/pengyuling/.conda/envs/yowo2/lib/python3.10/site-packages/coremltools/converters/mil/mil/input_type.py", line 163, in validate_inputs
raise ValueError(msg.format(name, var.name, input_type.type_str,
ValueError: Op "pred_reg.1" (op_type: gather) Input indices="anchor_idxs.1" expects tensor or scalar of dtype from type domain ['int32'] but got tensor[is1,fp32]

[TobyRoseman]

I would try updating the coremltools code to cast the indices to a dtype of int32.

Try adding the following line:
indices = mb.cast(x=indices, dtype="int32")
To coremltools/converters/mil/frontend/torch/ops.py in the index method right before the following line:
x = mb.gather(x=x, indices=indices, axis=axis, name=node.name)