File size: 724 Bytes
a76a14d
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
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);
}