Method __init__

python/dgl/optim/pytorch/sparse_optim.py:703–727

(
        self,
        params,
        lr,
        betas=(0.9, 0.999),
        eps=1e-08,
        use_uva=None,
        dtype=th.float32,
    )

    """

    def __init__(
        self,
        params,
        lr,
        betas=(0.9, 0.999),
        eps=1e-08,
        use_uva=None,
        dtype=th.float32,
    ):
        super(SparseAdam, self).__init__(params, lr)
        self._lr = lr
        self._beta1 = betas[0]
        self._beta2 = betas[1]
        self._eps = eps
        self._use_uva = use_uva
        self._nd_handle = {}
        self._is_using_uva = {}
        assert dtype in [th.float16, th.float32], (
            "Unsupported dtype {}. Valid choices are th.float32 "
            "and th.float32".format(dtype)
        )
        self._dtype = dtype

        # setup tensors for optimizer states
        self.setup(self._params)

    def _setup_uva(self, name, mem, power):
        self._is_using_uva[name] = True