use concat in sage then support batch training

xiang526 · Mar 30, 2020 · 91f52fc · 91f52fc
1 parent 3f147fe
commit 91f52fc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gcn.py b/gcn.py
@@ -143,7 +143,7 @@ def __init__(self, in_channels, out_channels, gcn_type, gcn_partition=None):
         else:
             if gcn_type == 'gat':
                 self.adj_available = False
-            if gcn_type in ['normal', 'cheb']:
+            if gcn_type in ['normal', 'cheb', 'graph', 'sage']:
                 self.batch_training = True
                 self.kwargs['node_dim'] = 1
             if gcn_type == 'cheb':