Function autograd::ops::argmax

pub fn argmax(x: &Tensor, axis: isize, keep_dim: bool) -> Tensor

Takes argmax along specified axis.