class RNN(nn.Module):

def __init__(self):
    super(RNN, self).__init__()
    self.rnn = nn.RNN(
        input_size=INPUT_SIZE,
        hidden_size=32, 
        num_layers=1,  
    )
    self.out = nn.Linear(32, 1)

def forward(self, x, h):
    
    out, h = self.rnn(x, h)
    prediction = self.out(out)
    return prediction, h


class RNN(nn.Module):
    def __init__(self):
    super(RNN, self).__init__()

    self.rnn = nn.RNN( 
        input_size=1,
        hidden_size=32,     
        num_layers=1,      
        batch_first=True,  
    )
    self.out = nn.Linear(32, 1)

def forward(self, x, h_state):  
    r_out, h_state = self.rnn(x, h_state)  
    outs = []   
    for time_step in range(r_out.size(1)):   
        outs.append(self.out(r_out[:, time_step, :]))
    return torch.stack(outs, dim=1), h_state