Internal change

PiperOrigin-RevId: 493705461

Author: learned_optimization authors

learned_optimization/research/hysteresis/truncated_pes_no_hyst.py

@@ -56,7 +56,7 @@ def init_worker_state(self, worker_weights: gradient_learner.WorkerWeights,
     return self.grad_est.init_worker_state(worker_weights, key)
 
   @profile.wrap()
-  def compute_gradient_estimate(
+  def compute_gradient_estimate(  # pytype: disable=signature-mismatch
       self,
       worker_weights: gradient_learner.WorkerWeights,
       key: PRNGKey,