diff --git a/tensorflow_probability/python/vi/optimization.py b/tensorflow_probability/python/vi/optimization.py
index <HASH>..<HASH> 100644
--- a/tensorflow_probability/python/vi/optimization.py
+++ b/tensorflow_probability/python/vi/optimization.py
@@ -135,10 +135,10 @@ def fit_surrogate_posterior(target_log_prob_fn,
       Default value: 'fit_surrogate_posterior'.
 
   Returns:
-    results: `Tensor` or nested structure of `Tensor`s, according to the
-      return type of `result_fn`. Each `Tensor` has an added leading dimension
-      of size `num_steps`, packing the trajectory of the result over the course
-      of the optimization.
+    results: `Tensor` or nested structure of `Tensor`s, according to the return
+      type of `trace_fn`. Each `Tensor` has an added leading dimension of size
+      `num_steps`, packing the trajectory of the result over the course of the
+      optimization.
 
   #### Examples
Fix a typo in the description of fit_surrogate_posterior