diff --git a/src/ErrorFunction/ErrorFunctions.cpp b/src/ErrorFunction/ErrorFunctions.cpp
index b0e12bfd0b3c7cce7de12c26cdbfb63dbfee88d7..e782851b664eaa42f27d3bbc72fef89d832fede4 100644
--- a/src/ErrorFunction/ErrorFunctions.cpp
+++ b/src/ErrorFunction/ErrorFunctions.cpp
@@ -56,10 +56,10 @@ namespace lib4neuro {
     void MSE::calculate_error_gradient(std::vector<double> &params, std::vector<double> &grad, double alpha, size_t batch) {
 
         size_t dim_out = this->ds->get_output_dim();
-        size_t n_elements = this->ds->get_n_elements();
+      
 
         std::vector<std::pair<std::vector<double>, std::vector<double>>> data = this->ds->get_random_data_batch(batch);
-
+		size_t n_elements = data.size();
         std::vector<double> error_derivative(dim_out);