2bff8339dd900a638be88441150166fd4c0fba19,tests/pytorch/test_nn.py,,test_set_trans,#,152

Before Change

    print(st_enc_0, st_enc_1, st_dec)

    // test/Ǘ: basic
    h0 = th.rand(g.number_of_nodes(), 50)
    h1 = st_enc_0(h0, g)
    assert h1.shape == h0.shape
    h1 = st_enc_1(h0, g)
    assert h1.shape == h0.shape
    h2 = st_dec(h1, g)
    assert h2.shape[0] == 200 and h2.dim() == 1

    // test/ǘ: batched graph
    g1 = dgl.DGLGraph(nx.path_graph(5))
    g2 = dgl.DGLGraph(nx.path_graph(10))
    bg = dgl.batch([g, g1, g2])
    h0 = th.rand(bg.number_of_nodes(), 50)
    h1 = st_enc_0(h0, bg)
    assert h1.shape == h0.shape
    h1 = st_enc_1(h0, bg)
    assert h1.shape == h0.shape

After Change

    st_enc_0 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, "sab")
    st_enc_1 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, "isab", 3)
    st_dec = nn.SetTransformerDecoder(50, 5, 10, 100, 2, 4)
    if F.gpu_ctx():
        st_enc_0.cuda()
        st_enc_1.cuda()
        st_dec.cuda()
    print(st_enc_0, st_enc_1, st_dec)

    // test/Ǘ: basic
    h0 = F.randn((g.number_of_nodes(), 50))

Instances