diff --git a/main_custom.py b/main_custom.py
index b7bba3e2cdd3f510761820dabc871ef18dbc9223..b1306505c4d41073b6f9a650333b544dc08087c4 100644
--- a/main_custom.py
+++ b/main_custom.py
@@ -220,7 +220,7 @@ def main(args):
     print('Cuda : ', torch.cuda.is_available())
 
     if args.forward == 'both':
-        data_train, data_val, _ = common_dataset.load_data(path_train=args.dataset_train,
+        data_train, data_val = common_dataset.load_data(path_train=args.dataset_train,
                                                                    path_val=args.dataset_val,
                                                                    path_test=args.dataset_test,
                                                                    batch_size=args.batch_size, length=args.seq_length, pad = False, convert=False, vocab='unmod')