Method _setup_uva

python/dgl/optim/pytorch/sparse_optim.py:729–733

(self, name, mem, power)

        self.setup(self._params)

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

    def setup(self, params):
        # We need to register a state sum for each embedding in the kvstore.