Method state_dict

python/dgl/optim/pytorch/sparse_optim.py:455–474

Return a copy of the whole optimizer states stored in CPU memory. If this is a multi-processing instance, the states will be returned in shared memory. If the underlying embedding is currently stored on multiple GPUs, all processes must call this method in the same order.

(self, **kwargs)

        self._clean_grad = True

    def state_dict(self, **kwargs):  # pylint: disable=unused-argument
        """Return a copy of the whole optimizer states stored in CPU memory.
        If this is a multi-processing instance, the states will be returned in
        shared memory. If the underlying embedding is currently stored on
        multiple GPUs, all processes must call this method in the same order.

        NOTE: This method must be called by all processes sharing the
        underlying embedding, or it may result in a deadlock.

        Returns
        -------
        dictionary of optimizer states
            The optimizer states stored in CPU memory.
        """
        return {
            "state": {
                emb.name: emb._all_get_optm_state() for emb in self._params
            },
            "param_groups": self.param_groups,
        }

    def load_state_dict(
        self, state_dict, **kwargs