Difficulty with checkpoint saving and loading (trainer+ FSDP accelerate)

Hi,
I am fine-tuning a 13B parameter using FSDP in Hugginface Trainer using accelerate. I want to load the best model at the end for which I have set the “load_best_model_at_end = True”, “save_total_limit=1” and “save_only_model=True”. The checkpointing is happening as expected, with older checkpoints being deleted automatically. However, while loading the best chekpoint at the end I am getting the error FileNotFoundError: [Errno 2] No such file or directory: 'output-dir/checkpoint-63/pytorch_model_fsdp.bin. I checked the checkpoint directory and found the checkpoints are being saved in shards as “pytorch_model-00001-of-00011.bin”.
How can I save the checkpoints in unsharded format so that it can be loaded, or enable loading the sharded checkpoint, either of which should resolve this issue.
@sgugger