tcsn_wty:

criterion = nn.L1Loss()


Here it is.
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable

# Parameters and DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100

class Model(nn.Module):
    # Our model
    
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc_a = nn.Linear(input_size, output_size)  # Unique A
        self.fc_b = nn.Linear(input_size, output_size)  # Unique B
        self.fc_out = nn.Linear(output_size, output_size)  # Shared layer
    
    def forward(self, input, select='A'):
        if select is 'A':
            input = self.fc_a(input)
        else:
            input = self.fc_b(input)
        output = self.fc_out(input)
        print("\tIn Model: output size", output.size())
        return output

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs

device = "cuda"
model = Model(input_size, output_size)
model = nn.DataParallel(model)  # For multi-gpu
model = model.module  # For multi-gpu
model.to(device)

optimizer_1 = optim.Adam([{'params': model.fc_a.parameters()},
                          {'params': model.fc_out.parameters()}],
                         lr=5e-5
                         )
optimizer_2 = optim.Adam([{'params': model.fc_b.parameters()},
                          {'params': model.fc_out.parameters()}],
                         lr=5e-5
                         )
criterion = nn.L1Loss()

batch_size = 16

print(model)

epochs = 10
for epoch in range(epochs):
    data_a, target_a = torch.randn(batch_size, data_size, input_size), torch.randn(batch_size, data_size, output_size)
    data_b, target_b = torch.randn(batch_size, data_size, input_size), torch.randn(batch_size, data_size, output_size)
    input_a, target_a = Variable(data_a.float()).to(device), Variable(target_a.float()).to(device)
    input_b, target_b = Variable(data_b.float()).to(device), Variable(target_b.float()).to(device)
    
    output_a = model(input_a, select='A')
    loss1 = criterion(output_a, target_a)
    loss1.backward()
    optimizer_1.step()
    print("Outside: output size", output_a.size())
    
    output_b = model(input_b, select='B')
    loss2 = criterion(output_b, target_b)
    loss2.backward()
    optimizer_2.step()
    print("Outside: output size", output_b.size())

Sorry for lacking of comments, please let me know if any explanation is needed.
By running this piece of code, I got the following kind of output:
        In Model: output size torch.Size([16, 100, 2])
Outside: output size torch.Size([16, 100, 2])
        In Model: output size torch.Size([16, 100, 2])
Outside: output size torch.Size([16, 100, 2])
        In Model: output size torch.Size([16, 100, 2])
Outside: output size torch.Size([16, 100, 2])

But I was expecting the following:
        In Model: output size torch.Size([8, 100, 2])
        In Model: output size torch.Size([8, 100, 2])
Outside: output size torch.Size([16, 100, 2])