@quanguet
Thank you for your reply.
I am not too sure what you mean by a detailed implementation. However, I missed out on giving my loss and optimizer function.
        self.model = Network(config).float()
        self.model = self.model.to(self.device)
        print(self.model)
        # Optimizer and Loss
        self.optimizer = optim.Adam(self.model.parameters(), lr=lr)
        self.pos_weight_factor = torch.tensor(120)
        # self.criterion = nn.BCELoss(reduction='none').to(self.device)
        self.criterion = nn.BCEWithLogitsLoss(pos_weight= self.pos_weight_factor)

What is the purpose of linear layer with hidden output? With BCEWithLogitsLoss I thought that I would need to put out as an input[predictions] to the loss = self.criterion(predictions.float(), target.float()). This is because BCEWIthLogisloss has intenal sigmoid function.
If you need any more information, please let me know