bef529dbb4eb10b1ec1cad14a6f630dd472dd73b,test/nn/conv/test_tag_conv.py,,test_tag_conv,#,5

Before Change

    conv.eval()
    jitconv.eval()
    jittedconv.eval()
    assert (torch.abs(conv(x, edge_index) -
            jitconv(x, edge_index)) < 1e-6).all().item()
    assert (torch.abs(conv(x, edge_index) -
            jittedconv(x, edge_index)) < 1e-6).all().item()

    assert conv(x, edge_index, edge_weight).size() == (num_nodes, out_channels)

After Change

    conv = TAGConv(in_channels, out_channels)
    assert conv.__repr__() == "TAGConv(16, 32, K=3)"
    out1 = conv(x, edge_index)
    assert out1.size() == (num_nodes, out_channels)
    out2 = conv(x, edge_index, edge_weight)
    assert out2.size() == (num_nodes, out_channels)

    jit_conv = conv.jittable(x=x, edge_index=edge_index)
    jit_conv = torch.jit.script(jit_conv)

Instances