diff --git a/tensorflow_graphics/math/optimizer/levenberg_marquardt.py b/tensorflow_graphics/math/optimizer/levenberg_marquardt.py
index 71cd5ed8e..cd0b8ff6e 100644
--- a/tensorflow_graphics/math/optimizer/levenberg_marquardt.py
+++ b/tensorflow_graphics/math/optimizer/levenberg_marquardt.py
@@ -78,7 +78,7 @@ def _compute_jacobian(values, variables, tape):
   values = tf.expand_dims(values, axis=-1)
   return values, jacobian
 
-
+@tf.function
 def minimize(residuals,
              variables,
              max_iterations,