Unexpected behavior with GatedGraphConv on batched graphs

[Theo]

I am having unexpected behavior while using the GatedGraphConv layer on batched graphs. Specifically, when testing with a batch of 2 identical graphs, the outputs for the first and second graph are not the same. Here is my code:

import dgl
import numpy as np
import torch as torch
from dgl.nn import GatedGraphConv
import networkx as nx

conv = GatedGraphConv(4, 1, 1, 3)
a = dgl.graph(([0, 1],[1, 2]))
a.ndata['node'] = torch.ones((3,1))
a.edata['edge'] = torch.Tensor([[0],[1]])

a = dgl.add_self_loop(a)
batch = dgl.batch([a,a])

print(conv(a,a.ndata['node'],a.edata['edge']))
print(conv(batch,batch.ndata['node'],batch.edata['edge']))

And the output is:
tensor([[0.9503],
[0.8060],
[0.8060]], grad_fn=)
tensor([[0.9503],
[0.8060],
[0.8060],
[0.9503],
[0.9753],
[0.8060]], grad_fn=)

Thank you for any help!

[Rhett-Ying]

Hi,

the issue you hit is caused by the incorrect shape of etypes you passed into forward function. pls replace it by a.edata['edge'] = torch.Tensor([[0],[1]]).squeeze().

BTW, the first argument of GatedGraphConv() should be input feature size, which should be 1 in your case, you passed in 4 instead.

please check the class/function signature in NN Modules (PyTorch) — DGL 0.7.0 documentation