Tornac:

The examples I’ve seen use classes derived from nn.Module which doesn’t fit my needs unfortunately.


I’m unsure why this wouldn’t fit your need. Could you explain your concerns a bit?
The general approach for your use case would be:

define modelA, modelB, and modelC

register all models in a parentModel

pass two inputs to parentModel.forward, call each branch using the corresponding data, concatenate the activations, and feed it to modelC


Here is an example:
class ParentModel(nn.Module):
    def __init__(self, modelA, modelB, modelC):
        super(ParentModel, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        self.modelC = modelC
        
    def forward(self, x1, x2):
        x1 = self.modelA(x1)
        x2 = self.modelB(x2)
        x = torch.cat((x1, x2), dim=1)
        x = self.modelC(x)
        return x
        
modelA = nn.Linear(10, 10)
modelB = nn.Linear(10, 10)
modelC = nn.Linear(20, 10)
parent = ParentModel(modelA, modelB, modelC)
x1 = torch.randn(1, 10)
x2 = torch.randn(1, 10)
out = parent(x1, x2)
print(out.shape)
> torch.Size([1, 10])

Instead of the nn.Linear modules for modelX you can also use nn.Sequential if you want.