Lightning Fabric: generic method to get the full state dict
Xynonners opened this issue · comments
Description & Motivation
Hi,
I would like a generic method of getting the full state dict on rank 0 for a lightning fabric wrapped module. Base lightning seems to have lightning_module_state_dict, but fabric only has save_checkpoint (which for my usecase doesn't work, as I do not want to save the entire state to a single file).
Pitch
See above.
Alternatives
None.
Additional context
No response
cc @Borda