Ensemble two features

[ptrblck]

You could replace the last linear layers with nn.Identity modules and create the ensemble as described in this post.