Fixed data bug

CRIPAC-DIG · Apr 13, 2022 · ccd032f · ccd032f
1 parent 68f9ca9
commit ccd032f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/new_data.py b/new_data.py
@@ -113,11 +113,11 @@ def generate_user(user, data, graph, item_max_length, user_max_length, train_pat
                 save_graphs(train_path+ '/' + str(user) + '/'+ str(user) + '_' + str(j) + '.bin', fin_graph,
                             {'user': torch.tensor([user]), 'target': torch.tensor([target]), 'u_alis':u_alis, 'last_alis': last_alis})
                 train_num += 1
-            if j == split_point-1-2:
+            if j == split_point - 1 - 1:
                 save_graphs(val_path + '/' + str(user) + '/' + str(user) + '_' + str(j) + '.bin', fin_graph,
                             {'user': torch.tensor([user]), 'target': torch.tensor([target]), 'u_alis': u_alis,
                              'last_alis': last_alis})
-            else:
+            if j == split_point - 1:
                 save_graphs(test_path + '/' + str(user) + '/' + str(user) + '_' + str(j) + '.bin', fin_graph,
                             {'user': torch.tensor([user]), 'target': torch.tensor([target]), 'u_alis':u_alis, 'last_alis': last_alis})
                 test_num += 1