如何在pytorch中收集每个观察值的预测？

import torch
from torch.utils.data import DataLoader

# 定义自定义回调函数
class PredictionCollector:
    def __init__(self):
        self.predictions = []

    def __call__(self, model, inputs, outputs):
        _, predicted = torch.max(outputs.data, 1)
        self.predictions.extend(predicted.tolist())

# 创建回调函数实例
collector = PredictionCollector()

# 加载数据集
dataset = YourDataset()
dataloader = DataLoader(dataset, batch_size=32)

# 创建模型实例
model = YourModel()

# 训练过程中注册回调函数
for inputs, labels in dataloader:
    outputs = model(inputs)
    # 其他训练步骤...

    # 调用回调函数收集预测值
    collector(model, inputs, outputs)

# 打印收集到的预测值
print(collector.predictions)