I am trying to train multioutput image classification (age and gender labels) but stuck in this problem. Here is my model. I have seen other results so far, unfortunately could not solve it.
class Net(nn.Module):
    
    def __init__(self):
        super(Net, self).__init__()
        
        self.conv_layer = nn.Sequential(
            
            # Conv Layer block 1
            nn.Conv2d(3, 64, 1, stride = 1),
            nn.SELU(inplace=True),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, 64, 7, stride = 1),
            nn.SELU(inplace=True),
            nn.MaxPool2d(3, 3),
            nn.BatchNorm2d(64),
            
            # Conv Layer block 2
            nn.Conv2d(64, 128, 1, stride = 1),
            nn.SELU(inplace=True),
            nn.BatchNorm2d(128),
            nn.Conv2d(128, 128, 5, stride = 1),
            nn.SELU(inplace=True),
            nn.MaxPool2d(3, 3),
            nn.BatchNorm2d(128),
            
            # Conv Layer block 3
            nn.Conv2d(128, 256, 1, stride = 1),
            nn.SELU(inplace=True),
            nn.BatchNorm2d(256),
            nn.Conv2d(256, 256, 3, stride = 2),
            nn.SELU(inplace=True),
            nn.BatchNorm2d(256),
            nn.Conv2d(256, 384, 1, stride = 1),
            nn.SELU(inplace=True),
            nn.MaxPool2d(2,2),
            nn.BatchNorm2d(384))
        
        self.fc_layer1 = nn.Sequential(
            nn.Linear(384*5*5, 1024),
            nn.SELU(inplace=True),
            nn.AlphaDropout(0.5),
            nn.Linear(1024, 512),
            nn.SELU(inplace=True),
            nn.AlphaDropout(0.5),            
            nn.Linear(512, 8)
        )
        
        self.fc_layer2 = nn.Sequential(
            nn.Linear(384*5*5, 1024),
            nn.SELU(inplace=True),
            nn.AlphaDropout(0.5),
            nn.Linear(1024, 512),
            nn.SELU(inplace=True),
            nn.AlphaDropout(0.5),            
            nn.Linear(512, 2)
        )
        
        
        self.apply(weights_init)
        
    def forward(self, x):
        # conv layers
        x = self.conv_layer(x)
        
        # flatten
        x = x.view(x.size(0), -1)
        
        label1 = self.fc_layer1(x)
        label2 = torch.sigmoid(self.fc_layer2(x))
        
        return {'label1': label1, 'label2': label2}