Joint Parameter Multivariate RNN

Hello, I am trying to extract joint features from H time series based on a single RNN. I have data tensors X of shape N x C x H x L, where N batch size, C is the number of features used as the input_size for the RNN, and L is the sequence length.

I’d like to use a single RNN to extract features from the H parallel time series. For this, I loop through the H rows (N x C x L) of the tensor X and feed them one after the other into the same RNN. The class CustomRNN below defines this RNN and takes care of reshaping the batch correctly (needs to be fed as (N x L x C) to the RNN. The CustomRNN maps from C → C hidden size and outputs a sequence of length L, again (with all the hidden states). The class JointRNN defines a single CustomRNN and loops through the rows of X, feeding them to the RNN and concatenating them to the final tensor. Is this the correct approach to extract joint RNN features from the individual time series?

class CustomRNN(nn.Module):

    def __init__(self,C): 
        super().__init__()
        self.rnn = nn.RNN(C,C,batch_first=True)
        return 
    
    def forward(self,x): 
        # input:  N x C x L
        # return: N x C x 1 x L
        return self.rnn(x.transpose(1,2))[0].transpose(1,2).unsqueeze(2)

class JointRNN(nn.Module): 
    
    def __init__(self,C,H,**kwargs): 
        super().__init__()
        self.H = H    
        self.rnn = CustomRNN(C)     
        return 
    
    def forward(self,x): 
        out = torch.tensor([],device=device)
        for i in range(self.H): 
            out = torch.cat([out,self.rnn(x[:,:,i,:])],dim=2)
        # return: N x C x H x L
        return out  

Thanks!