Fix 'Base' bug THUDM#28 THUDM#30

Aliang-CN · Oct 31, 2019 · 52b4ecd · 52b4ecd
1 parent b19c97c
commit 52b4ecd
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/main.py b/src/main.py
@@ -195,6 +195,10 @@ def train_model(network_data, feature_dic, log_name):
     train_pairs = generate_pairs(all_walks, vocab)
 
     edge_types = list(network_data.keys())
+    if edge_types[-1] != 'Base':
+        edge_types.sort()
+        edge_types.remove('Base')
+        edge_types.append('Base')
 
     num_nodes = len(index2word)
     edge_type_count = len(edge_types) - 1