[Distillation] save last chkpt as pytorch_model.bin

2019-09-09 19:53:35 +00:00
parent 84d346b687
commit a95ced6260
1 changed files with 4 additions and 1 deletions
--- 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,