Fix the type error. nn.module.get_variable returns a regular dict now.

PiperOrigin-RevId: 574546327
google · Oct 18, 2023 · d7121ed · d7121ed
1 parent 5dc9ded
commit d7121ed
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/vizier/_src/jax/stochastic_process_model.py b/vizier/_src/jax/stochastic_process_model.py
@@ -523,9 +523,7 @@ def posterior_predictive(
     )
     kwargs = cached_intermediates
     if not y_observed._nopadding_done:  # pylint: disable=protected-access
-      kwargs = kwargs.copy(
-          {'observations_is_missing': y_observed.is_missing[0]}
-      )
+      kwargs = kwargs | {'observations_is_missing': y_observed.is_missing[0]}
 
     predictive_index_points = tfpke.ContinuousAndCategoricalValues(
         continuous=x_predictive.continuous.padded_array,