fix eval

PiperOrigin-RevId: 612503759
google-research · Mar 4, 2024 · 828fdd9 · 828fdd9
1 parent 6ec92a8
commit 828fdd9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/t5x/eval.py b/t5x/eval.py
@@ -30,6 +30,7 @@
 os.environ['FLAX_LAZY_RNG'] = 'no'
 from absl import logging
 from clu import metric_writers
+import gin
 import jax
 import seqio
 from t5x import checkpoints
@@ -61,6 +62,7 @@ def __call__(
     ...
 
 
+@gin.configurable
 class InferenceEvaluator:
   """Runs evaluation of the model against a given SeqIo task."""