IndexError: index 4 is out of bounds for dimension 0 with size 4?

I am tying to write a NN in pytorch that learns to give for example: very simple A+B for two images.
Writing it, I got this error. my batch_size is 4.

class myNet(nn.Module):
    def __init__(self):
        super(myNet, self).__init__()
        
        self.fc1 = nn.Linear(3072, 3072)  # set up FC layer
        self.fc2 = nn.Linear(3072, 3072)  # set up the other FC layer
        self.fc3 = nn.Linear(3072, 3072)  # set up the other FC layer

    def forward(self, input1, input2):
        a = self.fc1(input1)
        print(a.size())
        b = self.fc2(input2)
        print(b.size())
        # now we can reshape `c` and `f` to 2D and concat them
        combined = [a[i]+b[i] for i in range(3072)]
        print(combined.size())
        out = self.fc3(combined)
        print(out.size())
        return out

and then:

num_epochs = 10
losses = []
batch_size= 4

for epoch in range(num_epochs):
    for i, (im1, im2, mid) in enumerate(trainloader):
        #im1= im1/255
        im1 = torch.from_numpy(np.array(im1, dtype='float32'))
        im1 = im1.to(device)
       
        
        #im2= im2/255
        im2 = torch.from_numpy(np.array(im2, dtype='float32'))
        im2 = im2.to(device)
        
        #mid= mid/255
        mid = torch.from_numpy(np.array(mid, dtype='float32'))
        mid = mid.to(device)
       
        # forwad pass
        outputs = model(im1, im2)
        #outputs = outputs.double()
        
        # loss
        loss = criterion(outputs, mid)
        losses.append(loss.item())

        # backward pass
        optimizer.zero_grad()
        loss.backward()
        
        # update parameters
        optimizer.step()
        
        # report
        if (i + 1) % 50 == 0:
            print('Epoch [%2d/%2d], Step [%3d/%3d], Loss: %.4f'
                  % (epoch + 1, num_epochs, i + 1, train_size // batch_size, loss.item()))

and running it, I get the mentioned error.
I know its relate to batches. But do not know how to fix it
Any ideas to fix it? or maybe another way to write this code is helpful.

>Solution :

On the following line:

combined = [a[i]+b[i] for i in range(3072)]

You are looping over the first dimension of a and b on [0, 3072[. Given that both tensors are shaped batch first, per PyTorch shape convention, you are going out of bound. What you are probably looking to do is loop over the second dimension instead which would correspond to your feature dimension.

combined = [a[:,i]+b[:,i] for i in range(3072)]

This is identical to using the torch.Tensor.__add__ operator:

combined = a + b 

Do note however that this is more of a addition fusion rather than a concatenation fusion.
Depending on your needs though, you might have to modify this piece of code with respect to how you actually want the two feature tensors to be merged.