I failed to save model when using torch.FSDP, what I tried to do is something like:
model = unwrap(self.model)
state_dict = model.state_dict()
But the process hangs forever, I have tried something like enable offload_to_cpu=True in this and this but nothing works.
I am using torch=2.1.0+cu117