Revert "[fix] oss dict load (#383)" (#384)

This reverts commit 8be9d93.

fairscale/optim/oss.py

@@ -391,16 +391,16 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
 
         # NOTE: PyTorch 1.5 does not index linearly but with the id(params) at saving time
         # we work around that here by using the fact that the params are ordered as in the param_groups
-        pytorch15_index_redirect = {k: i for i, k in enumerate(state_dict["state"].keys())}
 
-        for key, value in state_dict["state"].items():
-            param = self.index_to_param[pytorch15_index_redirect[key]]
+        for i_param, (key, value) in enumerate(state_dict["state"].items()):
+            param = self.index_to_param[i_param]
 
             # Populate the sharded optimizer state on the fly
             if self.param_to_rank[param] != self.rank:
                 state_dict["state"][key] = None
 
-            else:
+            if key in self.index_to_param:
+                param = self.index_to_param[i_param]
 
                 # Only add this state to the sharded optimizer if it owns this param
                 for pg in self.optim.param_groups: