Hello guys I am new to pytorch and I am trying to create the following neural network:

http://urban-computing.com/pdf/UbiComp2019JiZheng.pdf 1
The trick is that I have 5 copies of them and in the end I am joining them with a softmax function.
Am I doing this correctly? I do not have much experience with strange neural networks and how autograd works in them.
The input is ([batch_size, neural network number, input_state])
def net(torch.nn.Module):
def init(self, bases, state_space, action_space, n_hidden = 20):
super(policy_estimator,self).init()
    self.base11 = torch.nn.Linear(state_space,n_hidden)
    self.base12 = torch.nn.Linear(n_hidden, n_hidden)
    self.base13 = torch.nn.Linear(n_hidden, 1)

    self.base21 = torch.nn.Linear(state_space,n_hidden)
    self.base22 = torch.nn.Linear(n_hidden, n_hidden)
    self.base23 = torch.nn.Linear(n_hidden, 1)

    self.base31 = torch.nn.Linear(state_space,n_hidden)
    self.base32 = torch.nn.Linear(n_hidden, n_hidden)
    self.base33 = torch.nn.Linear(n_hidden, 1)

    self.base41 = torch.nn.Linear(state_space,n_hidden)
    self.base42 = torch.nn.Linear(n_hidden, n_hidden)
    self.base43 = torch.nn.Linear(n_hidden, 1)

    self.base51 = torch.nn.Linear(state_space,n_hidden)
    self.base52 = torch.nn.Linear(n_hidden, n_hidden)
    self.base53 = torch.nn.Linear(n_hidden, 1)

    # Activation functions
    self.activation_tanh = torch.nn.Tanh()
    self.activation_softmax = torch.nn.Softmax(dim=1)

def forward(self, x):

    out1  = self.activation_tanh(self.base11(x[:,0,:]))
    out1 = self.activation_tanh(self.base12(out1))
    out1 = self.base13(out1)

    out2 = self.activation_tanh(self.base21(x[:, 1, :]))
    out2 = self.activation_tanh(self.base22(out2))
    out2 = self.base23(out2)

    out3 = self.activation_tanh(self.base31(x[:, 2, :]))
    out3 = self.activation_tanh(self.base32(out3))
    out3 = self.base33(out3)

    out4 = self.activation_tanh(self.base41(x[:, 3, :]))
    out4 = self.activation_tanh(self.base42(out4))
    out4 = self.base43(out4)

    out5 = self.activation_tanh(self.base51(x[:, 4, :]))
    out5 = self.activation_tanh(self.base52(out5))
    out5 = self.base53(out5)

    score_torch=torch.stack([out1.squeeze(dim=0),out2.squeeze(dim=0),out3.squeeze(dim=0), out4.squeeze(dim=0), out5.squeeze(dim=0)], dim=1)
    action_probs = self.activation_softmax(score_torch)

    return action_probs

Thanks so much!