We can later move this to https://github.com/pytorch/examples or update this tutorial in https://github.com/pytorch/tutorials/blob/main/recipes_source/distributed_checkpoint_recipe.rst.

From what I learned, we should not keep the examples PT repo, but we still have a couples for checkpoint and other distributed componenets.

what does this do? should users rely on it, or should they write their own state function on their model?

This question is a bit loaded, I'm still thinking about the best UX around this item in particular. In general, if users are defining objects with custom Stateful behavior, they should define state_dict and load_state_dict on those objects.

For model/optim, which need to call get_state_dict and set_state_dict in state_dict and load_state_dict, the two options so far are:

1. Create a wrapper:

   class DOptim:
          def __init__(self, model, optim):
              self.model = model
              self.optim = optim
           ...
   

We're still evaluating whether it makes sense to include the wrapper as part of DCP since it can be a little tricky and could lead to a less then ideal UX

2. The other option, which I don't think is a bad one, is using the _patch methods as above. The patch methods are still in testing but I think it's pretty reasonable