Commit a95ced62 authored by VictorSanh's avatar VictorSanh
Browse files

[Distillation] save last chkpt as `pytorch_model.bin`

parent 84d346b6
...@@ -295,7 +295,10 @@ class Distiller: ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment