```
Note, that once load_state_dict_from_zero_checkpoint was run, the model will no longer be useable in the
DeepSpeed context of the same application. i.e. you will need to re-initialize the deepspeed engine, since
model.load_state_dict(state_dict) will remove all the DeepSpeed magic from it. So do this only at the very end
of the training.
Of course, you don't have to use class:~transformers.Trainer and you can adjust the examples above to your own
trainer.
If for some reason you want more refinement, you can also extract the fp32 state_dict of the weights and apply
these yourself as is shown in the following example:
```python
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)  # already on cpu
