This is my basic multitask learning model, and it has 2 tasks. Since there are only 2 tasks, maybe I can duplicate the code for each task as self.tower1 and self.tower2.

class Multitask_Network(nn.Module):
def __init__(self):

in_ch = 20
shared_layer_size = 10
tower_hidden_size = 10
output_size = 2

self.sharedlayer = nn.Sequential(
nn.Linear(in_ch, shared_layer_size),
nn.ReLU(),
nn.Dropout()
)
self.tower1 = nn.Sequential(
nn.Linear(shared_layer_size, tower_hidden_size),
nn.ReLU(),
nn.Dropout(),
nn.Linear(tower_hidden_size, output_size)
)
self.tower2 = nn.Sequential(
nn.Linear(shared_layer_size, tower_hidden_size),
nn.ReLU(),
nn.Dropout(),
nn.Linear(tower_hidden_size, output_size)
)

def forward(self, x):
x = self.sharedlayer(x)
out1 = self.tower1(x)
out2 = self.tower2(x)
return out1, out2

x = torch.ones((4, 20))
y = model(x)
y


But for my use case, I need to define the model's number of tasks during run time. I can do it with for-loop, but I think it is inefficient to run on GPU. I think it does the forward pass for task #1, then followed by task #2, to task #n, one after another.

How can I achieve this without a for loop in the forward pass?

class Multitask_Network(nn.Module):

in_ch = 20
shared_layer_size = 10
tower_hidden_size = 5
output_size = 2

self.sharedlayer = nn.Sequential(
nn.Linear(in_ch, shared_layer_size),
nn.ReLU(),
nn.Dropout()
)

self.conv_towers = []

tower = nn.Sequential(
nn.Linear(shared_layer_size, tower_hidden_size),
nn.ReLU(),
nn.Dropout(),
nn.Linear(tower_hidden_size, output_size)
)

self.conv_towers.append(
tower
)
self.conv_towers = nn.ModuleList(self.conv_towers)

def forward(self, x):
x = self.sharedlayer(x)

output = []