File size: 551 Bytes
0d06934 | 1 2 3 4 5 6 7 8 9 10 11 12 13 | """coremltools 9.0 workaround: aten::Int on a 1-element const array fails in _cast."""
import numpy as np
from coremltools.converters.mil.frontend.torch import ops as _ops
_orig_cast = _ops._cast
def _cast(context, node, dtype, dtype_name):
x = _ops._get_inputs(context, node)[0]
if x.val is not None and getattr(x.val, "shape", ()) != () and np.size(x.val) == 1:
context.add(_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name))
return
return _orig_cast(context, node, dtype, dtype_name)
_ops._cast = _cast
|