def load_state_dict(
        self, state_dict, **kwargs
    ):  # pylint: disable=unused-argument
        """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.

        Parameters
        ----------
        state_dict : dictionary of optimizer states
            The global states to pull values from.
        """
        for emb in self._params:
            emb._all_set_optm_state(state_dict["state"][emb.name])
        self._set_param_groups(state_dict["param_groups"])
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.

        Parameters
        ----------
        state_dict : dictionary of optimizer states
            The global states to pull values from.
        
python/dgl/optim/pytorch/sparse_optim.py