Many thanks for your prompt reply 
Here is my forward definition. The network is quite big, but it is a UNet-based for semantic segmentation using a shared encoder.
    def forward(self, x):
        features_x1 = self.encoder(x[:, : self.in_channels // 2, :, :])
        features_x2 = self.encoder(x[:, self.in_channels // 2 :, :, :])
        features = [torch.cat([x1, x2], 1) for x1, x2 in zip(features_x1, features_x2)]
        features = [self.res[i].to(x.device)(x) for i, x in enumerate(features)]

        decoder_output = self.decoder(*features)
        masks = self.segmentation_head(decoder_output)

        return masks