Trax Multi class Neural Net error with Eval_task

OS: Trax 1.3.1

$ pip freeze | grep trax
# trax==1.3.7

$ pip freeze | grep tensor
# mesh-tensorflow==0.1.18
tensor2tensor==1.15.7
tensorboard==2.4.1
tensorboard-plugin-wit==1.8.0
tensorflow==2.4.1
tensorflow-addons==0.12.1
tensorflow-datasets==4.2.0
tensorflow-estimator==2.4.0
tensorflow-gan==2.0.0
tensorflow-hub==0.11.0
tensorflow-metadata==0.27.0
tensorflow-probability==0.7.0
tensorflow-text==2.4.3

$ pip freeze | grep jax
# jax==0.2.9
jaxlib==0.1.59
$ python -V
# Python 3.7.6

# Steps to reproduce:
...

# Error logs:
from trax.supervised import training

batch_size = 16
rnd.seed(271)

train_task = training.TrainTask(
    labeled_data=train_generator(batch_size=batch_size, shuffle=True),
    loss_layer=tl.WeightedCategoryCrossEntropy(),
    optimizer=trax.optimizers.Adam(0.01),
    n_steps_per_checkpoint=10,
)

eval_task = training.EvalTask(
    labeled_data=val_generator(batch_size=batch_size, shuffle=True),
    metrics=[tl.WeightedCategoryCrossEntropy(), tl.Accuracy()],
)

model = classifier()

output_dir = '~/model/'
output_dir_expand = os.path.expanduser(output_dir)
print(output_dir_expand)

# GRADED FUNCTION: train_model
def train_model(classifier, train_task, eval_task, n_steps, output_dir):
    '''
    Input: 
        classifier - the model you are building
        train_task - Training taskeval
        eval_task - Evaluation task
        n_steps - the evaluation steps
        output_dir - folder to save your files
    Output:
        trainer -  trax trainer
    '''
### START CODE HERE (Replace instances of 'None' with your code) ###
    training_loop = training.Loop(
                                classifier, # The learning model
                                train_task, # The training task
                                eval_task = eval_task, # The evaluation task
                                output_dir = output_dir) # The output directory

    training_loop.run(n_steps = n_steps)
### END CODE HERE ###

    # Return the training_loop, since it has the model.
    return training_loop

training_loop = train_model(model, train_task, eval_task, 100, output_dir_expand)

Error code:
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-67-b0084f5b34fa> in <module>
----> 1 training_loop = train_model(model, train_task, eval_task, 100, output_dir_expand)

<ipython-input-66-c4ad3dc3f1c2> in train_model(classifier, train_task, eval_task, n_steps, output_dir)
     16                                 train_task, # The training task
     17                                 eval_task = eval_task, # The evaluation task
---> 18                                 output_dir = output_dir) # The output directory
     19 
     20     training_loop.run(n_steps = n_steps)

TypeError: __init__() got an unexpected keyword argument 'eval_task'