invalid argument 5: k not in range for dimension at /opt/conda/conda-bld/pytorch_1607370172916/work/aten/src/THC/generic/THCTensorTopK.cu:26
报错原因:我训练的模型是二分类,所以根本没有五个类
torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)
input:一个tensor数据
k:指明是得到前k个数据以及其index
dim: 指定在哪个维度上排序, 默认是最后一个维度
largest:如果为True,按照大到小排序; 如果为False,按照小到大排序
sorted:返回的结果按照顺序返回
out:可缺省,不要