Sum of specific indices in pytorch?

Question

I have a neural network that outputs a tensor of size 12. After applying some calculations to this tensor, I need to reduce it to size 8 by adding the first four pairs and turning those into 1 dimension.

[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] -> [3, 7, 11, 15, 9, 10, 11, 12]

Is there an operation like this in pytorch that would still allow me to apply gradients?

Answer 1

The gradient computation is still calculated so you do not need to worry about that. If you tensor is of shape (batch_size, 12), you can do the following:

import torch 
output = torch.Tensor([
    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 
    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
])
output = torch.cat((
    output[:, 0].unsqueeze(1) + output[:, 1].unsqueeze(1),
    output[:, 2].unsqueeze(1) + output[:, 3].unsqueeze(1),
    output[:, 4].unsqueeze(1) + output[:, 5].unsqueeze(1),
    output[:, 6].unsqueeze(1) + output[:, 7].unsqueeze(1),
    output[:, 8:]),
dim=1)
print(output)


>>> tensor([[ 3.,  7., 11., 15.,  9., 10., 11., 12.],
        [ 3.,  7., 11., 15.,  9., 10., 11., 12.]])

By the way, why not using nn.Linear()?