I’m not sure I understand the code completely, but it looks like you are storing some activations in self.encoder_outs. Are these activations for debugging only or is this list part of the model?
In the latter case, how did you define self.encoder_outs and how would you like to reduce it?
I guess each replica will initialize its own self.encoder_outs, so that you end up having N lists.