From c2b4ce02b5e7d5df85810fa3fad6eb04d1d25cf6 Mon Sep 17 00:00:00 2001 From: V-MoE Authors Date: Wed, 16 Aug 2023 00:13:00 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 557378061 --- vmoe/evaluate/evaluator.py | 2 +- vmoe/projects/adversarial_attacks/attacks.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vmoe/evaluate/evaluator.py b/vmoe/evaluate/evaluator.py index a6d1d61..bd3b637 100644 --- a/vmoe/evaluate/evaluator.py +++ b/vmoe/evaluate/evaluator.py @@ -138,7 +138,7 @@ def _make_callback_fn(cls, *, apply_fn, loss_fn, label_pred_fn, datasets, # Note: We create the eval_step_pjit here to avoid multiple compilation # steps. If the shapes of inputs/outputs for all datasets is the same, this # will be only compiled once. - eval_state_dtype_struct = EvalState( + eval_state_dtype_struct = EvalState( # pytype: disable=wrong-arg-types # dataclass_transform num=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32), sum_correct=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32), sum_loss=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32), diff --git a/vmoe/projects/adversarial_attacks/attacks.py b/vmoe/projects/adversarial_attacks/attacks.py index b5269b0..a81bd8a 100644 --- a/vmoe/projects/adversarial_attacks/attacks.py +++ b/vmoe/projects/adversarial_attacks/attacks.py @@ -20,7 +20,7 @@ import jax.numpy as jnp -class AttackState(flax.struct.PyTreeNode): +class AttackState(flax.struct.PyTreeNode): # pytype: disable=invalid-function-definition # dataclass_transform """State used to aggregate attack statistics over several batches of images. Attributes: