如何在pytorch中获取混淆矩阵？

import torch
import numpy as np
from sklearn.metrics import confusion_matrix

# 定义模型
model = YourModel()

# 定义数据加载器
data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)

model.eval()

all_predictions = []
all_labels = []

with torch.no_grad():
    for inputs, labels in data_loader:
        # 将数据传入模型进行预测
        outputs = model(inputs)
        _, predictions = torch.max(outputs, 1)
        
        # 将预测结果和真实标签添加到列表中
        all_predictions.extend(predictions.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

confusion_mat = confusion_matrix(all_labels, all_predictions)