class Dscnvo(nn.Module):
    def __init__(self):
        super(Dscnvo, self).__init__()
        self.conv1   = nn.Conv2d(6, 64, 7, 2, 3)
        self.conv2   = nn.Conv2d(64, 128, 5, 2, 2)
        self.conv3   = nn.Conv2d(128, 256, 5, 2, 2)
        self.conv3_1 = nn.Conv2d(256, 256, 3, 1, 1)
        self.conv4   = nn.Conv2d(256, 512, 3, 2, 1)
        self.conv4_1 = nn.Conv2d(512, 512, 3, 1, 1)
        self.conv5   = nn.Conv2d(512, 512, 3, 2, 1)
        self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 1)
        self.conv6   = nn.Conv2d(512, 1024, 3, 2, 1)
        self.conv6_1 = nn.Conv2d(1024, 1024, 2, 2, 0)
        self.maxpool = nn.MaxPool2d((2, 2), stride = (2, 2))
        self.fc1     = nn.Linear(5120, 4096)
        self.fc2     = nn.Linear(4096, 1024)
        self.fc3     = nn.Linear(1024, 128)
        self.fc4     = nn.Linear(128, 6)
        self.init_weights()
    
    def forward(self, img):
        x = F.relu(self.conv1(img))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv3_1(x))
        x = F.relu(self.conv4(x))
        x = F.relu(self.conv4_1(x))
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv5_1(x))
        x = F.relu(self.conv6(x))
        x = self.conv6_1(x)
        x = self.maxpool(x)
        x = x.view(-1, 5120)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return x

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d or nn.Linear):
                nn.init.xavier_normal_(m.weight)

This is my above model input stack of images - dim [batch_size, 6, 384, 1280] output [ batch_size, 6]
def forward_pass(batch):
    imgf, imgr, tf, af = batch
    imgf, imgr, tf, af = imgf.to(device),imgr.to(device), tf.to(device), af.to(device)
    #imgr = torch.cat((imgf[:, 3:6, :, :],imgf[:, 0:3, :, :]), 1)
    tr   = tf*-1
    ar   = af*-1
    outf = model(imgf)
    outr = model(imgr)
    loss1 = (alpha1*(criterion(outf[:, 0:3].unsqueeze(1), tf)) + beta1*(criterion(outf[:, 3:6].unsqueeze(1), af)))
    loss2 = (alpha2*(criterion(outr[:, 0:3].unsqueeze(1), tr)) + beta2*(criterion(outr[:, 3:6].unsqueeze(1), ar)))
    loss  = loss1 + loss2
    return loss, outf, tf, af

for epoch in range(1, num_epochs+1):
    running_loss = 0.0
    for i, batch in enumerate(trainloader):
        #forward
        loss, outf, tf, af = forward_pass(batch)

        #backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item()