| """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 | |