PyTorch:将钩子添加到模型以保存中间层输出将返回两次要素

import torch
import torch.nn as nn

# 定义一个简单的模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(128 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 定义一个钩子函数
def hook_fn(module, input, output):
    # 保存中间层的输出
    intermediate_output = output.clone().detach()
    torch.save(intermediate_output, 'intermediate_output.pt')

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

# 添加钩子到指定层
model.conv1.register_forward_hook(hook_fn)

# 使用模型进行前向传播
input = torch.randn(1, 3, 32, 32)
output = model(input)

# 保存中间层的输出
torch.save(output, 'output.pt')