Fix LayerNorm Implementation

src/mlpack/methods/ann/layer/layer_norm.hpp

@@ -184,6 +184,9 @@ class LayerNorm
 
   //! Locally-stored normalized input.
   OutputDataType normalized;
+
+  //! Locally-stored input with 0 mean.
+  OutputDataType inputMean;
 }; // class LayerNorm
 
 } // namespace ann

src/mlpack/methods/ann/layer/layer_norm_impl.hpp

@@ -63,7 +63,7 @@ void LayerNorm<InputDataType, OutputDataType>::Forward(
 
   // Normalize the input.
   output = input.each_row() - mean;
-
+  inputMean = output;
   output.each_row() /= arma::sqrt(variance + eps);
 
   // Reused in the backward and gradient step.
@@ -79,7 +79,6 @@ template<typename eT>
 void LayerNorm<InputDataType, OutputDataType>::Backward(
     const arma::Mat<eT>&& input, arma::Mat<eT>&& gy, arma::Mat<eT>&& g)
 {
-  const arma::mat inputMean = input.each_row() - mean;
   const arma::mat stdInv = 1.0 / arma::sqrt(variance + eps);
 
   // dl / dxhat