From a95ced6260494f108df72f0f9ffe9c60498365ad Mon Sep 17 00:00:00 2001 From: VictorSanh Date: Mon, 9 Sep 2019 19:53:35 +0000 Subject: [PATCH] [Distillation] save last chkpt as `pytorch_model.bin` --- examples/distillation/distiller.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/distillation/distiller.py b/examples/distillation/distiller.py index ed710a2bee..93135e292c 100644 --- a/examples/distillation/distiller.py +++ b/examples/distillation/distiller.py @@ -295,7 +295,10 @@ class Distiller: if self.is_master: logger.info(f'--- Ending epoch {self.epoch}/{self.params.n_epoch-1}') self.end_epoch() - if self.is_master: logger.info('Training is finished') + if self.is_master: + logger.info(f'Save very last checkpoint as `pytorch_model.bin`.') + self.save_checkpoint(checkpoint_name=f'pytorch_model.bin') + logger.info('Training is finished') def step(self, input_ids: torch.tensor,