Я получаю сообщение об ошибке при использовании нейронной сети trax

Question

python3

Я пытаюсь обучить нейронную сеть с помощью trax, это для курса NLP, нейронная сеть предназначена для анализа настроений. Я использую блокнот jupyter, это ошибка, в которой она сокращена:

LayerError Traceback (последний вызов последним) в ----> 1 training_loop = train_model(модель, train_task, eval_task, 100, output_dir_expand)

in train_model(classifier, train_task, eval_task, n_steps, output_dir)19 output_dir = output_dir) # Каталог вывода 20---> 21 training_loop.run(n_steps = n_steps)22 ### КОНЕЦ КОДА ЗДЕСЬ ###23

TypeError: mul получил несовместимые формы для трансляции: (8, 2), (16, 2). это код:

      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.CrossEntropyLoss(),
    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.CrossEntropyLoss(), tl.Accuracy()],
)

model = classifier()



def train_model(classifier, train_task, eval_task, n_steps, output_dir):
    
    #Input: 
       # classifier - the model you are building
        #train_task - Training task
        #eval_task - Evaluation task
        #n_steps - the evaluation steps
        #output_dir - folder to save your files
    #Output:
        #trainer -  trax trainer


    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)


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