torch简单构造最后一层输出为前面层输出的乘积

我的网络有N层，我想要最后一层简单地将输出构造为前面层的输出的乘积，如下所示。假设我的第N-1层具有3维输出，最后一层具有2维输出。最后一层输出[1] = previous_output[1] * previous_output[2]，最后一层输出[2] = previous_output[3]。我想在这个2D的last_output上使用CrossEntropyCriterion。现在我通过构建N-1网络模型并获取以下最后一个输出来实现这一点：

last_output = torch.ones(previous_output:size()[1], 2)
last_output[{{}, {1}}] = torch.cmul(previous_output[{{},{1}}], previous_output[{{},{2}}])
last_output[{{}, {2}}] = previous_output[{{}, {3}}]

因此，最后一层不在模型中，当简单地运行

loss = criterion:forward(last_output, target)

会发生错误。 是否有什么技巧来实现这个目的？