diff --git a/tensorflow_probability/python/optimizer/bfgs_test.py b/tensorflow_probability/python/optimizer/bfgs_test.py
index 0077794ca6..e7cc1d1d9e 100644
--- a/tensorflow_probability/python/optimizer/bfgs_test.py
+++ b/tensorflow_probability/python/optimizer/bfgs_test.py
@@ -14,8 +14,6 @@
 # ============================================================================
 """Tests for the unconstrained BFGS optimizer."""
 
-import functools
-
 from absl.testing import parameterized
 import numpy as np
 from scipy.stats import special_ortho_group
@@ -30,7 +28,6 @@
 
 
 def _make_val_and_grad_fn(value_fn):
-  @functools.wraps(value_fn)
   def val_and_grad(x):
     return gradient.value_and_gradient(value_fn, x)
   return val_and_grad
diff --git a/tensorflow_probability/python/optimizer/lbfgs_test.py b/tensorflow_probability/python/optimizer/lbfgs_test.py
index 77ed64f4d9..891ac47b53 100644
--- a/tensorflow_probability/python/optimizer/lbfgs_test.py
+++ b/tensorflow_probability/python/optimizer/lbfgs_test.py
@@ -14,8 +14,6 @@
 # ============================================================================
 """Tests for the unconstrained L-BFGS optimizer."""
 
-import functools
-
 from absl.testing import parameterized
 import numpy as np
 from scipy.stats import special_ortho_group
@@ -30,7 +28,6 @@
 
 
 def _make_val_and_grad_fn(value_fn):
-  @functools.wraps(value_fn)
   def val_and_grad(x):
     return gradient.value_and_gradient(value_fn, x)
   return val_and_grad