torch.Tensor.argmax

Tensor.argmax(dim=None, keepdim=False) → LongTensor

See torch.argmax()