Support RelGCN in dataset class
@@ -79,7 +79,7 @@ class RelGCN(chainer.Chain):
self.input_type = input_type
self.scale_adj = scale_adj
- def __call__(self, x, adj):
+ def __call__(self, batch):
"""
Args:
x: (batchsize, num_nodes, in_channels)
@@ -87,6 +87,8 @@ class RelGCN(chainer.Chain):
Returns: (batchsize, hidden_channels)
"""
+ x, adj = batch.x, batch.adj
+
if x.dtype == self.xp.int32:
assert self.input_type == 'int'
else:
