diff --git a/conmech/solvers/optimization/optimization.py b/conmech/solvers/optimization/optimization.py
index 76e4e857..8c3eed0d 100644
--- a/conmech/solvers/optimization/optimization.py
+++ b/conmech/solvers/optimization/optimization.py
@@ -34,7 +34,7 @@
 from conmech.solvers.solver_methods import (
     make_cost_functional,
     make_equation,
-    make_cost_functional_subgradient,
+    make_subgradient,
 )
 
 
@@ -62,11 +62,9 @@ def __init__(
             variable_dimension=statement.dimension_out,
             problem_dimension=statement.dimension_in,
         )
-        if hasattr(contact_law, "subderivative_normal_direction"):
-            self.subgradient = make_cost_functional_subgradient(
-                djn=contact_law.subderivative_normal_direction,  # TODO
-                djt=None,
-                dh_functional=None,
+        if hasattr(contact_law, "subderivative_normal_direction"):  # TODO
+            self.subgradient = make_subgradient(
+                djn=contact_law.subderivative_normal_direction,
             )
         if isinstance(statement, WaveStatement):
             if isinstance(contact_law, InteriorContactLaw):
diff --git a/conmech/solvers/solver_methods.py b/conmech/solvers/solver_methods.py
index f28f15f5..2fcafe77 100644
--- a/conmech/solvers/solver_methods.py
+++ b/conmech/solvers/solver_methods.py
@@ -161,9 +161,8 @@ def contact_cost(length, normal, normal_bound, tangential, tangential_bound):
     def contact_cost_functional(
         var, var_old, static_displacement, nodes, contact_boundary, contact_normals, dt
     ):
-        offset = len(var) // problem_dimension
-
         cost = 0.0
+
         # pylint: disable=not-an-iterable
         for ei in numba.prange(len(contact_boundary)):
             edge = contact_boundary[ei]
@@ -188,10 +187,6 @@ def contact_cost_functional(
                 static_displacement_mean - static_displacement_normal * normal_vector
             )
 
-            for node_id in edge:
-                if node_id >= offset:
-                    continue
-
             cost += contact_cost(
                 nph.length(edge, nodes),
                 normal_condition(vm_normal, static_displacement_normal, dt),
@@ -226,18 +221,16 @@ def cost_functional(
     return cost_functional
 
 
-def make_cost_functional_subgradient(
-    djn: Callable, djt: Optional[Callable] = None, dh_functional: Optional[Callable] = None
+def make_subgradient(
+    djn: Callable,
+    problem_dimension=2,
 ):
     djn = njit(djn)
-    djt = njit(djt)
-    dh_functional = njit(dh_functional)
-    DIMENSION = 2
 
     @numba.njit()
-    def contact_subgradient(u_vector, u_vector_old, nodes, contact_boundary, contact_normals):
+    def contact_subgradient(u_vector, nodes, contact_boundary):
         cost = np.zeros_like(u_vector)
-        offset = len(u_vector) // DIMENSION
+        offset = len(u_vector) // problem_dimension
 
         for edge in contact_boundary:
             n_id_0 = edge[0]
@@ -254,7 +247,7 @@ def contact_subgradient(u_vector, u_vector_old, nodes, contact_boundary, contact
                 cost[n_id_1 + offset] = cost[n_id_1]
         return cost
 
-    # pylint: disable=unused-argument # 'dt'
+    # pylint: disable=unused-argument # takes the same args as cost_functional
     @numba.njit()
     def subgradient(
         var,
@@ -268,7 +261,7 @@ def subgradient(
         base_integrals,
         dt,
     ):
-        dj = contact_subgradient(var, var_old, nodes, contact_boundary, contact_normals)
+        dj = contact_subgradient(var, nodes, contact_boundary)
         result = np.dot(lhs, var) - rhs + dj
         result = result.ravel()
         return result