Save custom objects in the state for each process

As of now custom registered objects are not saved per process as seen here. Is there a way to save a separate checkpoint and load for each of the register custom objects per process/rank in a distributed setting?

Alternatively, do you suggest that we should gather everything before saving and then distribute during loading?