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: