diff --git a/touvlo/utils.py b/touvlo/utils.py
index fdf47db..ca8531c 100644
--- a/touvlo/utils.py
+++ b/touvlo/utils.py
@@ -35,7 +35,7 @@ def g_grad(x):
 
 
 def gradient_descent(X, y, grad, initial_theta,
-                     alpha, num_iters, _lambda=None):
+                     alpha, num_iters, **kwargs):
     """This function performs parameter optimization via gradient descent.
 
     :param X: Features' dataset plus bias column.
@@ -56,22 +56,12 @@ def gradient_descent(X, y, grad, initial_theta,
     :param num_iters: Number of times the optimization will be performed.
     :type num_iters: int
 
-    :param _lambda: Weight of the penalty term.
-    :type _lambda: float
-
     :returns: Optimized model parameters.
     :rtype: numpy.array
     """
-    if _lambda is not None:
-        theta = copy(initial_theta)
-
-        for _ in range(num_iters):
-            theta = theta - alpha * grad(theta, X, y, _lambda)
-
-    else:
-        theta = copy(initial_theta)
-        for _ in range(num_iters):
-            theta = theta - alpha * grad(theta, X, y)
+    theta = copy(initial_theta)
+    for _ in range(num_iters):
+        theta = theta - alpha * grad(X, y, theta, **kwargs)
 
     return theta
 

Generalize the function utils.gradient_descent 
The function `gradient_descent` should only be aware of the parameters  `X, y, grad, initial_theta,alpha, num_iters`. Any further parameter should be fed to the `grad` function with no impact in the signature of the function `gradient_descent`.

TL; DR: Edit the function `gradient_descent` in the module `utils` to allow a gradient function (grad) of any signature.
diff --git a/tests/test_utils.py b/tests/test_utils.py
index d695ead..5adeb9b 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -4,7 +4,7 @@ import pytest
 from numpy import array, cos, sin, exp
 from numpy.testing import assert_allclose
 
-from touvlo.utils import numerical_grad, g_grad
+from touvlo.utils import numerical_grad, g_grad, gradient_descent
 
 
 class TestLogisticRegression:
@@ -82,3 +82,45 @@ class TestLogisticRegression:
         assert_allclose(g_grad(z),
                         [0.196612, 0.235004, 0.25, 0.235004, 0.196612],
                         rtol=0, atol=0.001, equal_nan=False)
+
+    def test_gradient_descent1(self, err):
+        def grad(X, y, theta):
+            m = len(y)
+            grad = (1 / m) * (X.T).dot(X.dot(theta) - y)
+            return grad
+
+        X = array([[0, 1, 2], [-1, 5, 3], [2, 0, 1]])
+        initial_theta = array([[0], [0], [0]])
+        y = array([[0.3], [1.2], [0.5]])
+        num_iters = 3
+        alpha = 1
+
+        assert_allclose(array([[-46.415], [276.248], [192.204]]),
+                        gradient_descent(X, y, grad, initial_theta,
+                                         alpha, num_iters),
+                        rtol=0, atol=0.001, equal_nan=False)
+
+    def test_gradient_descent2(self, err):
+        def grad(X, y, theta, schleem, plumbus, wubba, lubba):
+            m = len(y)
+            grad = (schleem / (m * wubba))
+            grad = grad * (X.T).dot(X.dot(theta) - y)
+            grad = grad + plumbus / (2 * lubba)
+            return grad
+
+        X = array([[0, 1, 2], [-1, 5, 3], [2, 0, 1]])
+        initial_theta = array([[0], [0], [0]])
+        y = array([[0.3], [1.2], [0.5]])
+        num_iters = 5
+        plumbus = 0.8
+        schleem = 0.6
+        wubba = 3.4
+        lubba = 2.7
+        alpha = 0.01
+
+        assert_allclose(array([[-0.0078777], [0.0106179], [0.0060865]]),
+                        gradient_descent(X, y, grad, initial_theta,
+                                         alpha, num_iters, lubba=lubba,
+                                         schleem=schleem, wubba=wubba,
+                                         plumbus=plumbus),
+                        rtol=0, atol=0.001, equal_nan=False)

pip install -e .[dev]
parse_log_pytest
requirements.txt
pytest
apt-get update apt-get install -y gcc
3.9
requirements.txt dev-requirements.txt
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
build==1.2.2.post1
check-manifest==0.50
cmake==4.0.0
exceptiongroup==1.2.2
flake8==7.2.0
flake8-import-order==0.18.2
importlib_metadata==8.6.1
iniconfig==2.1.0
mccabe==0.7.0
numpy==2.0.2
packaging==24.2
pbr==6.1.1
pluggy==1.5.0
pycodestyle==2.13.0
pyflakes==3.3.2
pyproject_hooks==1.2.0
pytest==8.3.5
tomli==2.2.1
-e git+https://github.com/Benardi/touvlo.git@0e1ee71ebec9a09370f928dd96d0ae11f3bba404#egg=touvlo
zipp==3.21.0

name: touvlo
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=3.0.16=h5eee18b_0
  - pip=25.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - setuptools=75.8.0=py39h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - build==1.2.2.post1
      - check-manifest==0.50
      - cmake==4.0.0
      - exceptiongroup==1.2.2
      - flake8==7.2.0
      - flake8-import-order==0.18.2
      - importlib-metadata==8.6.1
      - iniconfig==2.1.0
      - mccabe==0.7.0
      - numpy==2.0.2
      - packaging==24.2
      - pbr==6.1.1
      - pluggy==1.5.0
      - pycodestyle==2.13.0
      - pyflakes==3.3.2
      - pyproject-hooks==1.2.0
      - pytest==8.3.5
      - tomli==2.2.1
      - zipp==3.21.0
prefix: /opt/conda/envs/touvlo

tests/test_utils.py::TestLogisticRegression::test_gradient_descent1 tests/test_utils.py::TestLogisticRegression::test_gradient_descent2
tests/test_utils.py::TestLogisticRegression::test_numeric_grad_1 tests/test_utils.py::TestLogisticRegression::test_numeric_grad_2 tests/test_utils.py::TestLogisticRegression::test_numeric_grad_3 tests/test_utils.py::TestLogisticRegression::test_numeric_grad_4 tests/test_utils.py::TestLogisticRegression::test_numeric_grad_5 tests/test_utils.py::TestLogisticRegression::test_sigmoid_gradient
touvlo/utils.py
touvlo/utils.py