program(1.0) [buildInfo = dict, tensor>({{"coremlc-component-MIL", "5.33.5"}, {"coremlc-version", "1877.40.3"}, {"coremltools-component-torch", "2.1.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "7.2"}})] { func main(tensor logits) { tensor var_2 = const()[name = tensor("op_2"), val = tensor(-1)]; tensor var_3 = const()[name = tensor("op_3"), val = tensor(false)]; tensor argmax = reduce_argmax(axis = var_2, keep_dims = var_3, x = logits)[name = tensor("op_4_cast_fp16")]; } -> (argmax); }