Pytorch のnn.Linearの出力次元数を合わせたい

Question

以下のshapeの入力用のtensorがあります

torch.Size([4, 15232]) (preds_)

これをnn.Linearに通して次元を894まで減らしたいです。

# このshapeに
torch.Size([4, 896])

以下のコードでやったのですが、エラーがでてうまく行きません。
解決策を教えていただけないでしょうか？
ご教授お願いします。

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        INPUT_FEATURES = 4*15232
        OUTPUT = 896
        self.fc1 = nn.Linear(INPUT_FEATURES, OUTPUT)

    def forward(self, x):
        x = self.fc1(x)
        return x


fc = Net()
fc.to(device)

print(preds_.shape)
>>> torch.Size([4, 15232]) 
preds_ = fc(preds_)

# エラー
File "/usr/local/lib/python3.7/site-packages/torch/nn/functional.py", line 1690, in linear
　　　ret = torch.addmm(bias, input, weight.t())
RuntimeError: mat1 dim 1 must match mat2 dim 0

Answer 1

model.fc = nn.Linear(INPUT_FEATURES, OUTPUT)