Collects the state dict and dump to disk.
(trainer: transformers.Trainer, output_dir: str)
| 102 | ) |
| 103 | |
| 104 | def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str): |
| 105 | """Collects the state dict and dump to disk.""" |
| 106 | state_dict = trainer.model.state_dict() |
| 107 | if trainer.args.should_save: |
| 108 | cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()} |
| 109 | del state_dict |
| 110 | trainer._save(output_dir, state_dict=cpu_state_dict) # noqa |
| 111 | |
| 112 | |
| 113 | def smart_tokenizer_and_embedding_resize( |