[Distillation] save last chkpt as pytorch_model.bin

This commit is contained in:
VictorSanh
2019-09-09 19:53:35 +00:00
parent 84d346b687
commit a95ced6260

View File

@@ -295,7 +295,10 @@ class Distiller:
if self.is_master: logger.info(f'--- Ending epoch {self.epoch}/{self.params.n_epoch-1}') if self.is_master: logger.info(f'--- Ending epoch {self.epoch}/{self.params.n_epoch-1}')
self.end_epoch() 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, def step(self,
input_ids: torch.tensor, input_ids: torch.tensor,