This is a technical proof of concept prior to experimenting with an MoE RNN.
Essentially, I want to run two RNNs in parallel on the same input sequence. The caveat is that the number of parallel RNNs to run is unknown - thus the for loop.
class pRNN(nn.Module):

    def __init__(self, input_size=0, hidden_size=0, num_layers=1, bidirectional=False, dropout=0, subunit_count=1):

        super().__init__()

        self.subunit_count = subunit_count

        subunit_size = math.ceil(hidden_size / subunit_count)

        self.rnn = []
        hidden_size_remaining = hidden_size
        for i in range(self.subunit_count):
            self.rnn.append(nn.GRU(input_size=input_size, hidden_size=min(hidden_size_remaining, subunit_size), num_layers=num_layers, bidirectional=bidirectional, dropout=dropout))

        self.rnn = nn.ModuleList(self.rnn)

    def forward(self, x):

        out = None
        for i in range(self.subunit_count):
            if out is None:
                out, hidden = self.rnn[i](x)
            else:
                out2 = self.rnn[i](x)
                out = torch.cat((out, out2[0]), dim=-1)
                hidden = torch.cat((hidden, out2[1]), dim=-1)

        return out, hidden

Unfortunately, this is maybe the fifth failed attempt to come up with a technique that runs in parallel - no matter what I try it runs sequentially (including reducing the already minimal batch size of 128 by the number of RNNs). I am running on a single V100.
Any suggestions are appreciated.