Can't load optimizer state due to `state_steps`

Hi, I recently upgraded to PyTorch 1.12 and have had issues with loading a saved optimizer state using FSDP here and the issue seems something that is addressed in comments here -
https://github.com/facebookresearch/fairscale/blob/4975b05e89aaa29923b72c23b7b0f45118e4252f/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L2436

From what I understand, Adam's `step` state changed into a singleton tensor and when I call `gather_full_optim_state_dict()` this `step` is converted to an int.

Sample saving dict code:
```
model = FSDP(model, ...)
# call on all ranks
optim_state = model.gather_full_optim_state_dict(optimizer)
if rank == 0:
# save only on rank 0
checkpoint = {
'optimizer': optim_state,
...
}
torch.save(checkpoint)
```

Now when I load this optim state dict back - I do the following:
```
model = FSDP(model, ...)
torch.distributed.barrier()
# on all ranks
checkpoint = torch.load(snapshot_name)
curr_opt_state_dict = checkpoint["optimizer"]
optim_shard_dict = model.get_shard_from_optim_state_dict(curr_opt_state_dict)
optimizer.load_state_dict(optim_shard_dict)
```

This always fails the assertion in the Adam code - https://github.com/pytorch/pytorch/blob/master/torch/optim/adamw.py#L204 because I imagine the step was converted to an `int` within FSDP and Adam expects it to be a singleton tensor.

My question is am I saving the state dict correctly? Do I need to call `optimizer.state_dict()` on top of `model.gather_full_optim_state_dict()`?

A workaround I'm using to get things to bypass the assertion is to convert the `ints` back to singleton tensors in the adamw function however that does not seem safe. Any thoughts?

Apologies if my understanding is incorrect, I followed some of the discussion here - https://github.com/facebookresearch/fairscale/issues/776 for the state_dict saving logic.

Reproduce the save/load path using FSDP's gather_full_optim_state_dict() and get_shard_from_optim_state_dict(), then inspect the referenced FSDP code near line 2436 and the AdamW assertion near line 204. Compare the behavior with issue #776 and determine the expected representation for state_steps; done means the saved optimizer state loads without the assertion failure and the behavior is covered by a regression test.

Written by the indexing model from the issue text.