How to use trax to evaluate accuracy with all samples in EvalTask at each checkpoint?

Question

I am in the process of learning trax. I would like to implement a simple MLP in trax, so I translated the tensorflow version of the "MLP for Multiclass Classification" code at
https://machinelearningmastery.com/tensorflow-tutorial-deep-learning-with-tf-keras/
to a trax version:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import numpy as np
from numpy import argmax
import trax
from trax import layers as tl
from pandas import read_csv
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from trax.supervised import training

# load the dataset
path = 'https://raw.githubusercontent.com/jbrownlee/Datasets/master/iris.csv'
# dataset contains 150 entries of four elements long arrays and
# 150 long array of integer of 0,1,2
df = read_csv(path, header=None)
# split into input and output columns
X, y = df.values[:, :-1], df.values[:, -1]
# ensure all data are floating point values
X = X.astype('float32')
# encode strings to integer
y = LabelEncoder().fit_transform(y)
# split into train and test datasets, first 100 is train, last 50 is test
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33)
print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)
# determine the number of input features
n_features = X_train.shape[1]

mlp = tl.Serial(
    tl.Dense(10),
    tl.Relu(),
    tl.Dense(8),
    tl.Relu(),
    tl.Dense(3),
    tl.LogSoftmax()
)

train_task = training.TrainTask(
    labeled_data = zip(X_train, y_train),
    loss_layer = tl.CategoryCrossEntropy(),
    optimizer = trax.optimizers.Adam(0.01),
    sample_batch = (np.array([5.1,3.5,1.4,0.2], dtype=np.float32), np.int64(0)),
    n_steps_per_checkpoint = 25
)

eval_task = training.EvalTask(
    labeled_data = zip(X_test, y_test),
    metrics=[tl.CategoryCrossEntropy(), tl.CategoryAccuracy()],
    sample_batch = (np.array([5.1,3.5,1.4,0.2], dtype=np.float32), np.int64(0)),
    n_eval_batches=10
)

output_dir = os.path.expanduser('~/output_dir/')
os.system("rm -rf %s" % output_dir)
training_loop = training.Loop(mlp, train_task, eval_tasks=[eval_task], output_dir=output_dir)
training_loop.run(100)

row = np.array([5.1,3.5,1.4,0.2])
yhat = np.exp(mlp(row))
print('Predicted: %s (class=%d)' % (yhat, argmax(yhat)))

It seems to work but the Accuracy evalution part is not what I wanted. With the above code, Accuracy evaluation is done for every 25 training samples as well as the first sample. As we have 100 training samples, 4+1 evaluations are done and each time it uses 10 samples in the test set. If I change n_steps_per_checkpoint to 20, it will crash with StopIteration exception as now 5+1 evaluations are done and I need 60 samples in the test set when I only have 50.

I think what I wanted to do is to use all 50 test samples for evaluation at each checkpoint. Is it possible to do this without changing the trax source code? Or am I wrong in using all 50 samples for evaluation at each checkpoint? If so, what is the right way for evaluation? Thank you very much in advance.

My environment
tensorflow version: 2.10.0
jax version: 0.3.23
trax version: 1.4.1