DGLerror: DGLGraphsendではなく、DGLGraph.apply_edgesとsend_and_recvを代用する手法

Question

以下のサイトを参考にしてDGLライブラリを用いて、グラフ向け深層学習を実装しようとしたところ、以下のようなエラーが出ました。
https://qiita.com/K-1/items/62cd7fe04b80868ded8e

コードはここからダウンロードしました。
https://gist.github.com/k1ochiai/cd0279ca79dd74e91a2b5e1187928adb

利用環境は、
anaconda
python version3.6.12
UI: Jupyter Notebook
バックエンド: PyTorch
ライブラリ：DGL(Deep Graph Library), networkx

#コードの内容

import dgl
import torch

def build_sample_graph():
    g = dgl.DGLGraph()
    g.add_nodes(5)
    edge_list = [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (2, 3)]
    src, dst = tuple(zip(*edge_list))
    g.add_edges(src, dst)
    g.add_edges(dst, src)

    return g

G = build_sample_graph()
print('We have %d nodes.' % G.number_of_nodes())
print('We have %d edges.' % G.number_of_edges())

import networkx as nx
nx_G = G.to_networkx().to_undirected()
pos = nx.kamada_kawai_layout(nx_G)
nx.draw(nx_G, pos,with_labels=True,  node_color=[[.7, .7, .7]])

import torch
import torch.nn as nn
import torch.nn.functional as F

def gcn_message(edges):
    print("gcn_message:",edges.src['h'])
    return {'msg' : edges.src['h']}

def gcn_reduce(nodes):
    print("gcn_reduce:",nodes.mailbox['msg'])
    return {'h' : torch.sum(nodes.mailbox['msg'], dim=1)}

class GCNLayer(nn.Module):
    def __init__(self, in_feats, out_feats):
        super(GCNLayer, self).__init__()
        self.linear = nn.Linear(in_feats, out_feats)
 
    def forward(self, g, inputs):
        #フォワードパスを定義
        g.ndata['h'] = inputs
        g.send(g.edges(), gcn_message)
        g.recv(g.nodes(), gcn_reduce)
        h = g.ndata.pop('h')
        return self.linear(h)

class GCN(nn.Module):
    def __init__(self, in_feats, num_classes):
        super(GCN, self).__init__()
        self.gcn1 = GCNLayer(in_feats, num_classes)

    def forward(self, g, inputs):
        h = self.gcn1(g, inputs)
        h = torch.relu(h)

        return h

net = GCN(5, 2)
inputs = torch.eye(5)
labeled_nodes = torch.tensor([0, 1, 2, 3, 4])
labels = torch.tensor([0, 0, 1, 1, 0])

optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
all_logits = []
G.set_n_initializer(dgl.init.zero_initializer)

for epoch in range(3):    
    logits = net(G, inputs)
    all_logits.append(logits.detach())
    logp = F.log_softmax(logits, 1)
    loss = F.nll_loss(logp[labeled_nodes], labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print('Epoch %d | Loss: %.4f' % (epoch, loss.item()))

#バグの発生箇所（DGLerror)

---------------------------------------------------------------------------
DGLError                                  Traceback (most recent call last)
<ipython-input-20-adea207aebaa> in <module>
      4 
      5 for epoch in range(3):
----> 6     logits = net(G, inputs)
      7     all_logits.append(logits.detach())
      8     logp = F.log_softmax(logits, 1)

~/anaconda/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

<ipython-input-18-70c26494cadc> in forward(self, g, inputs)
      5 
      6     def forward(self, g, inputs):
----> 7         h = self.gcn1(g, inputs)
      8         h = torch.relu(h)
      9 

~/anaconda/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

<ipython-input-17-a96f9714ed3a> in forward(self, g, inputs)
     18     def forward(self, g, inputs):
     19         g.ndata['h'] = inputs
---> 20         g.send(g.edges(), gcn_message)
     21         g.recv(g.nodes(), gcn_reduce)
     22         h = g.ndata.pop('h')

~/anaconda/envs/pytorch/lib/python3.6/site-packages/dgl/heterograph.py in send(self, edges, message_func, etype)
   5556         DEPRECATE: please use send_and_recv, update_all.
   5557         """
-> 5558         raise DGLError('DGLGraph.send is deprecated. As a replacement, use DGLGraph.apply_edges\n'
   5559                        ' API to compute messages as edge data. Then use DGLGraph.send_and_recv\n'
   5560                        ' and set the message function as dgl.function.copy_e to conduct message\n'

DGLError: DGLGraph.send is deprecated. As a replacement, use DGLGraph.apply_edges
 API to compute messages as edge data. Then use DGLGraph.send_and_recv
 and set the message function as dgl.function.copy_e to conduct message
 aggregation.

と上記のように、DGLerrorが出ました。
DGLGraph.sendではなく、DGLGraph.apply_edgesとDGLGraph.send_and_recvを用いればいいと書かれてありますが、以上のコードをどのように書き換えれば良いのかがわかりません。ご教授願えませんでしょうか？