93fab2e53dbd812c8cb08fc7a789ec0ace9e8768,test/nn/conv/test_gcn_conv.py,,test_gcn_conv,#,5

Before Change

                                torch.Size([num_nodes, in_channels]))
    conv(x, edge_index).size() == (num_nodes, out_channels)
    jitcls = conv.jittable(x=x, edge_index=edge_index)
    jitconv = jitcls(in_channels, out_channels)
    jitconv.load_state_dict(conv.state_dict())
    jittedconv = torch.jit.script(jitconv)
    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()

    conv = GCNConv(in_channels, out_channels, cached=True)
    conv(x, edge_index).size() == (num_nodes, out_channels)
    conv(x, edge_index).size() == (num_nodes, out_channels)
    jitcls = conv.jittable(x=x, edge_index=edge_index)
    jitconv = jitcls(in_channels, out_channels, cached=True)
    jitconv.load_state_dict(conv.state_dict())
    jittedconv = torch.jit.script(jitconv)
    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()

After Change

    jit_conv = conv.jittable(x=x, edge_index=edge_index)
    jit_conv = torch.jit.script(jit_conv)
    jit_out = jit_conv(x, edge_index)
    assert jit_out.tolist() == out.tolist()

    conv = GCNConv(in_channels, out_channels, cached=True)
    out = conv(x, edge_index)

Instances