|
program(1.0) |
|
[buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"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<ios16>(tensor<fp16, [1, 64, 32000]> logits) { |
|
tensor<int32, []> var_2 = const()[name = tensor<string, []>("op_2"), val = tensor<int32, []>(-1)]; |
|
tensor<bool, []> var_3 = const()[name = tensor<string, []>("op_3"), val = tensor<bool, []>(false)]; |
|
tensor<int32, [1, 64]> argmax = reduce_argmax(axis = var_2, keep_dims = var_3, x = logits)[name = tensor<string, []>("op_4_cast_fp16")]; |
|
} -> (argmax); |
|
} |