how to use torch.distributed.checkpoint to save and load state dict

sparse_arch is a part of my model.

<img width="721" height="698" alt="Image" src="https://github.com/user-attachments/assets/cb35959b-418e-4ff4-8e12-4524528cbad2" />

<img width="1439" height="684" alt="Image" src="https://github.com/user-attachments/assets/d008966e-e2d2-404d-bcda-bce3e3285eed" />