tensor二维张量在行上去最大值，argmax

如果你想在一个二维张量（比如一个矩阵）的行上找到最大值，并返回这个最大值所在的列的索引，可以使用PyTorch中的`argmax`函数。例如，假设我们有一个2x3的矩阵`x`：

import torch
x = torch.tensor([[1, 2, 3], [4, 5, 6]])

我们可以使用`argmax`函数来找到每行的最大值所在的列：

max_indices = torch.argmax(x, dim=1)

这将返回一个一维张量，包含两个元素，每个元素是对应行的最大值所在的列的索引。在这个例子中，`max_indices`应该是一个长度为2的一维张量，其中第一个元素是2（因为第一行的最大值是3，它在第二列），第二个元素是2（因为第二行的最大值是6，它在第三列）。