Pytorch 2D Feature Tensor 1D output Одинаковые веса по измерению

Я обрабатываю набор данных временных рядов с n временными шагами, m функциями и k объектами. В результате мой вектор признаков имеет форму (n,k,m), в то время как форма моих целей (n,m)

Я хочу предсказать цели для каждого временного шага и объекта, но с одинаковыми весами для каждого объекта. Также моя функция потерь выглядит так.

average_loss = loss_func(prediction, labels)
sum_loss = loss_func(sum(prediction), sum(labels))
loss = loss_weight * average_loss + (1-loss_weight) * sum_loss

Мой план состоит не только в том, чтобы убедиться, что я предсказываю каждый элемент как можно лучше, но также в том, чтобы была предсказана сумма всех элементов. loss_weights — константа.

В настоящее время я делаю такое уродливое решение:

features = local_batch.squeeze(dim = 0)
labels = torch.unsqueeze(local_labels.squeeze(dim = 0), 1)
prediction = net(features)

Я устанавливаю свой размер партии = 1. И сжимаю его, чтобы сделать k объектов моей партией.

Моя сеть выглядит так:

def __init__(self, n_feature, n_hidden, n_output):
    super(Net, self).__init__()
    self.hidden = torch.nn.Linear(n_feature, n_hidden)   # hidden layer
    self.predict = torch.nn.Linear(n_hidden, n_output)   # output layer

def forward(self, x):
    x = F.relu(self.hidden(x))      # activation function for hidden layer
    x = self.predict(x)             # linear output
    return x

Как мне убедиться, что я делаю разумную свертку по измерению объекта, чтобы сохранить одинаковые веса для всех объектов, не привязываясь к размеру партии = 1? Кроме того, как мне добиться той же функции потерь, где я вычисляю потерю суммы прогноза по сравнению с целевой суммой для любой метки времени?