@ptrblck, thank you for the toy example. How would you change the script if modelB 's input was taken from one of the layers output of modelA. And the ensemble output has 2 heads.In the modification below, fc1 of modelA is ran twice (not efficient):
class MyModelA(nn.Module):
    def __init__(self):
        super(MyModelA, self).__init__()
        self.fc1 = nn.Linear(20, 10)
        self.fc2 = nn.Linear(10, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    

class MyModelB(nn.Module):
    def __init__(self):
        super(MyModelB, self).__init__()
        self.fc1 = nn.Linear(10, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        return x

class MyEnsemble(nn.Module):
    def __init__(self, modelA, modelB):
        super(MyEnsemble, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        
    def forward(self, x1, x2):
        head1 = self.modelA(x1)
        head2 = self.modelB(self.modelA.fc1(x1))
        x = torch.cat((head1, head2), dim=1)
        return x

Thanks!