Week 2 Lab - Missing Save Command

For the section ’ Fine-Tune the Model with the Preprocessed Dataset", there is a missing save command.

It currently reads:


output_dir = f'./dialogue-summary-training-{str(int(time.time()))}'

training_args = TrainingArguments(
    output_dir=output_dir,
    learning_rate=1e-5,
    num_train_epochs=1,
    weight_decay=0.01,
    logging_steps=1,
    max_steps=1
)

trainer = Trainer(
    model=original_model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation']
)

trainer.train()

The following statement is necessary to save the model:

trainer.save_model()

3 Likes

Hi Vikram, and welcome to the community! Another learner mentioned something similar. This might have been skipped because a pretrained model was used for the next sections. But it would indeed be nice to have it for completeness. We’ll consider it and update the notebook if needed. Thanks!