Load the optimizer states. This method must be called by all processes sharing the underlying embedding with identical :attr:`state_dict`. NOTE: This method must be called by all processes sharing the underlying embedding, or it may result in a deadlock. Par
(
self, state_dict, **kwargs
)
| 474 | } |
| 475 | |
| 476 | def load_state_dict( |
| 477 | self, state_dict, **kwargs |
| 478 | ): # pylint: disable=unused-argument |
| 479 | """Load the optimizer states. This method must be called by all |
| 480 | processes sharing the underlying embedding with identical |
| 481 | :attr:`state_dict`. |
| 482 | |
| 483 | NOTE: This method must be called by all processes sharing the |
| 484 | underlying embedding, or it may result in a deadlock. |
| 485 | |
| 486 | Parameters |
| 487 | ---------- |
| 488 | state_dict : dictionary of optimizer states |
| 489 | The global states to pull values from. |
| 490 | """ |
| 491 | for emb in self._params: |
| 492 | emb._all_set_optm_state(state_dict["state"][emb.name]) |
| 493 | self._set_param_groups(state_dict["param_groups"]) |
| 494 | |
| 495 | @property |
| 496 | @abstractmethod |