def _set_n_repr(self, ntid, u, data):
        """Internal API to set node features.

        `data` is a dictionary from the feature name to feature tensor. Each tensor
        is of shape (B, D1, D2, ...), where B is the number of nodes to be updated,
        and (D1, D2, ...) be the shape of the node representation tensor. The
        length of the given node ids must match B (i.e, len(u) == B).

        All updates will be done out of place to work with autograd.

        Parameters
        ----------
        ntid : int
            Node type id.
        u : node, container or tensor
            The node(s).
        data : dict of tensor
            Node representation.
        """
        if is_all(u):
            num_nodes = self._graph.num_nodes(ntid)
        else:
            u = utils.prepare_tensor(self, u, "u")
            num_nodes = len(u)
        for key, val in data.items():
            nfeats = F.shape(val)[0]
            if nfeats != num_nodes:
                raise DGLError(
                    "Expect number of features to match number of nodes (len(u))."
                    " Got %d and %d instead." % (nfeats, num_nodes)
                )
            if F.context(val) != self.device:
                raise DGLError(
                    'Cannot assign node feature "{}" on device {} to a graph on'
                    " device {}. Call DGLGraph.to() to copy the graph to the"
                    " same device.".format(key, F.context(val), self.device)
                )
            # To prevent users from doing things like:
            #
            #     g.pin_memory_()
            #     g.ndata['x'] = torch.randn(...)
            #     sg = g.sample_neighbors(torch.LongTensor([...]).cuda())
            #     sg.ndata['x']    # Becomes a CPU tensor even if sg is on GPU due to lazy slicing
            if (
                self.is_pinned()
                and F.context(val) == "cpu"
                and not F.is_pinned(val)
            ):
                raise DGLError(
                    "Pinned graph requires the node data to be pinned as well. "
                    "Please pin the node data before assignment."
                )

        if is_all(u):
            self._node_frames[ntid].update(data)
        else:
            self._node_frames[ntid].update_row(u, data)
Internal API to set node features.

        `data` is a dictionary from the feature name to feature tensor. Each tensor
        is of shape (B, D1, D2, ...), where B is the number of nodes to be updated,
        and (D1, D2, ...) be the shape of the node representation tensor. The
        length of the given node ids must match B (i.e, len(u) == B).

        All updates will be done out of place to work with autograd.

        Parameters
        ----------
        ntid : int
            Node type id.
        u : node, container or tensor
            The node(s).
        data : dict of tensor
            Node representation.
        
python/dgl/heterograph.py