diff --git a/kratos/spaces/ublas_space.h b/kratos/spaces/ublas_space.h
index 70857c2f9bd8..a41981cde8d1 100644
--- a/kratos/spaces/ublas_space.h
+++ b/kratos/spaces/ublas_space.h
@@ -708,10 +708,16 @@ class UblasSpace
         const SCALING_DIAGONAL ScalingDiagonal = SCALING_DIAGONAL::NO_SCALING
         )
     {
+        // The system size
         const std::size_t system_size = rA.size1();
 
-        const auto& Avalues = rA.value_data();
-        const auto& Arow_indices = rA.index1_data();
+        // The matrix data
+        auto& r_Avalues = rA.value_data();
+        const auto& r_Arow_indices = rA.index1_data();
+        const auto& r_Acol_indices = rA.index2_data();
+
+        // Define the iterators
+        const auto it_Acol_indices_begin = r_Acol_indices.begin();
 
         // Define  zero value tolerance
         const double zero_tolerance = std::numeric_limits<double>::epsilon();
@@ -723,18 +729,29 @@ class UblasSpace
         IndexPartition(system_size).for_each([&](std::size_t Index){
             bool empty = true;
 
-            const std::size_t col_begin = Arow_indices[Index];
-            const std::size_t col_end = Arow_indices[Index + 1];
+            const std::size_t col_begin = r_Arow_indices[Index];
+            const std::size_t col_end = r_Arow_indices[Index + 1];
 
             for (std::size_t j = col_begin; j < col_end; ++j) {
-                if(std::abs(Avalues[j]) > zero_tolerance) {
+                if(std::abs(r_Avalues[j]) > zero_tolerance) {
                     empty = false;
                     break;
                 }
             }
 
             if(empty) {
-                rA(Index, Index) = scale_factor;
+                const auto it_Acol_indices_row_begin = it_Acol_indices_begin + col_begin;
+                const auto it_Acol_indices_row_end = it_Acol_indices_begin + col_end;
+
+                const auto lower = std::lower_bound(it_Acol_indices_row_begin, it_Acol_indices_row_end, Index);
+                const auto upper = std::upper_bound(it_Acol_indices_row_begin, it_Acol_indices_row_end, Index);
+
+                if (lower != upper) { // Index was found
+                    r_Avalues[std::distance(it_Acol_indices_begin, lower)] = scale_factor;
+                } else {
+                    KRATOS_DEBUG_ERROR << "Diagonal term (" << Index << ", " << Index << ") is not defined in the system matrix" << std::endl;
+                    KRATOS_WARNING("UblasSpace") << "Diagonal term (" << Index << ", " << Index << ") is not defined in the system matrix" << std::endl;
+                }
                 rb[Index] = 0.0;
             }
         });