PyTorch torch.argmin 函数
Pytorch torch 参考手册torch.argmin 是 PyTorch 中用于返回沿维度最小值的索引的函数。
函数定义
torch.argmin(input, dim, keepdim)
使用示例
实例
import torch
x = torch.tensor([[1, 3, 2], [4, 1, 3]])
print("全局最小索引:", torch.argmin(x))
print("dim=1 最小索引:", torch.argmin(x, dim=1))
x = torch.tensor([[1, 3, 2], [4, 1, 3]])
print("全局最小索引:", torch.argmin(x))
print("dim=1 最小索引:", torch.argmin(x, dim=1))
输出结果为：
全局最小索引: tensor(0) dim=1 最小索引: tensor([0, 1])