Modernize @jaxtyped code

PiperOrigin-RevId: 590708836
google · Dec 13, 2023 · e64328f · e64328f
1 parent 4b14b9c
commit e64328f
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/vizier/_src/algorithms/designers/scalarization.py b/vizier/_src/algorithms/designers/scalarization.py
@@ -35,8 +35,7 @@ class Scalarization(abc.ABC, eqx.Module):
   # Weights shape should be broadcastable with objectives when called.
   weights: jt.Float[jax.Array, '#Obj'] = eqx.field(converter=jnp.asarray)
 
-  @jt.jaxtyped
-  @typeguard.typechecked
+  @jt.jaxtyped(typechecker=typeguard.typechecked)
   def __call__(
       self, objectives: jt.Float[jax.Array, '*Batch Obj']
   ) -> jt.Float[jax.Array, '*Batch']: