PyTorch中计算Fisher Metric遇NNGeometry维度报错的解决与替代方案

我有一个用于求解矩阵方程Ax=b的简易PyTorch模型（A为3×3矩阵，x和b为3×1列向量），模型仅包含一个线性层，训练代码如下：

import torch
import torch.nn as nn
import torch.optim as optim

class Net(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Net, self).__init__()

        self.linear = nn.Linear(input_dim, output_dim, bias=False)

    def forward(self, x):
        out = self.linear(x)
        return out


# Define the training data
A = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.],
                  [7., 8., 9.]])

b = torch.tensor([[52.],
                  [124.],
                  [196.]])

# Define the model and the optimizer
model = Net(input_dim=9, output_dim=3)
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Train the model
for epoch in range(2000):
    optimizer.zero_grad()
    y_pred = model(A.reshape(9))
    print(A@y_pred[:3])
    loss = nn.MSELoss(reduction='sum')(A@y_pred.view((3,1)), b)
    loss.backward()
    optimizer.step()

# Evaluate the model
with torch.no_grad():
    y_pred = model(A.reshape(9))
    print("Solution:\n", y_pred)

训练完成后，我希望计算该模型的Fisher Metric。尝试使用NNGeometry包进行计算，该包需要DataLoader，因此我构建了仅包含训练矩阵A的数据集和DataLoader，代码如下：

class TrivialDataset(Dataset):
    def __init__(self):
        self.data = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]).reshape(1,9)
    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return len(self.data)

# Create the DataLoader
dataset = TrivialDataset()
loader = DataLoader(dataset, batch_size=1)

最后尝试生成Fisher信息矩阵（FIM），代码如下：

from nngeometry.metrics import FIM
from nngeometry.object import PMatDense
fisher_metric = FIM(model, loader, n_output=1, variant='regression', representation=PMatDense, device='cpu')

但运行时出现错误：

RuntimeError: shape '[9, 1]' is invalid for input of size 3

我发现问题源于维度转换操作，但NNGeometry应该支持输入维度大于输出维度的模型（例如分类任务），请问如何规避该报错？是否有NNGeometry的优质替代方案？

内容的提问来源于stack exchange，提问作者notastringtheorist