FSDP model can't save checkpoint

1

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

2

I ran into same issue. Have you solved it yet?