File size: 599 Bytes
deaf348 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | """coremltools 9 + numpy 2: the torch frontend's _cast folds 1-element constants with int(array([n])), which numpy 2 rejects."""
import numpy as np
from coremltools.converters.mil.frontend.torch import ops as _ops
_cast0 = _ops._cast
def _cast(context, node, dtype, dtype_name):
x = _ops._get_inputs(context, node, expected=1)[0]
if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
context.add(_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
return
_cast0(context, node, dtype, dtype_name)
_ops._cast = _cast
|