def forward(self, x: torch.Tensor) -> ttorch.Tensor:
        """
        **Overview**:
            The computation graph of the linear layer with popart mechanism, which outputs both the output and the normalized output of the layer.
        """
        # Execute the linear layer computation $$y=Wx+b$$, note here we use expand and broadcast to add bias.
        normalized_output = x.mm(self.weight.t())
        normalized_output += self.bias.unsqueeze(0).expand_as(normalized_output)
        # Unnormalize the output for more convenient usage.
        with torch.no_grad():
            output = normalized_output * self.sigma + self.mu

        return ttorch.as_tensor({'output': output, 'normalized_output': normalized_output})

        **Overview**:
            The computation graph of the linear layer with popart mechanism, which outputs both the output and the normalized output of the layer.
        
chapter4_reward/popart.py