logits = torch.norm(out, dim=-1) pred = torch.eye(10).to(device).index_select(dim=0, index=torch.argmax(logits, dim=1))

这是一个基于 PyTorch 的分类模型的预测代码，其中 logits 是模型输出的未经 softmax 处理的预测值，torch.norm 函数用于计算每个样本的预测值向量的 L2 范数，torch.argmax 函数用于找到每个样本预测值向量中最大值的下标，最后使用 torch.eye 函数生成一个 10 维的 one-hot 向量，再使用 index_select 函数根据最大值下标选择对应的 one-hot 向量作为预测结果。