@smth, are you saying this is the preferred method in the case of independent subnets?
class Net(nn.Module):
   def __init__(self, num, embedding_dim):
       super().__init__()

       # list of independent DCNN nets
       self.sub_nets = nn.ModuleList([subnet(output_size = embedding_dim)
                                for i in range(num)])
       
       ...
   
   def forward(self, x):
       
       # forward on the expert nets
       subset_outs = torch.stack([model(x) for model in self.sub_nets])
       
       ...