Home
last modified time | relevance | path

Searched refs:argmax_kernel_cuda_impl (Results 1 – 1 of 1) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DReduceArgMaxKernel.cu21 void argmax_kernel_cuda_impl(TensorIterator& iter) { in argmax_kernel_cuda_impl() function
34 argmax_kernel_cuda_impl<at::Half, float>(iter); in argmax_kernel_cuda()
36 argmax_kernel_cuda_impl<at::BFloat16, float>(iter); in argmax_kernel_cuda()
39 argmax_kernel_cuda_impl<scalar_t>(iter); in argmax_kernel_cuda()