I would like to train an ensemble of same network architectures. Currently, I’m defining it this way.
def make_net(network_specs):
    # return nn.Module

class Ensemble(nn.Module):
    def __init__(self, network_specs, ensemble_size):
        super().__init__()
        self.model = nn.ModuleList([make_net(network_specs) for _ in range(ensemble_size)])
        
    def forward(self, x):
        return torch.stack([self.model[i](x[i]) for i in range(ensemble_size)])

However, backprop of the stack operator doesn’t seem to be parallelized (in the same GPU). The GPU utilization is very low, around 15%.
Thinking that dynamic graph might be the cause (similar thread here 5), I’ve recently tried using @torch.jit. The performance still doesn’t improve.
What am I doing wrong here? How can I improve the performance of my ensemble model?
Thanks.