diff --git a/pytorch_classifier.py b/pytorch_classifier.py
index 02310d86d0a170918529561954b91bd74ecab1f7..808c4fe56187e7747c6fdeebe0120d0e315e1730 100644
--- a/pytorch_classifier.py
+++ b/pytorch_classifier.py
@@ -121,7 +121,13 @@ class TransformerClassifier(nn.Module):
 model = TransformerClassifier()
 
 
-def train(model, train_dataloader, dev_dict_dataloader, test_dict_sentences, epochs, specific_results):
+def train(model, 
+          train_dataloader, 
+          dev_dict_dataloader, 
+          test_dict_sentences, 
+          test_dict_dataloader,
+          epochs, 
+          specific_results):
 
     device = torch.device("cuda" if args.use_cuda else "cpu")
 
@@ -153,6 +159,7 @@ def train(model, train_dataloader, dev_dict_dataloader, test_dict_sentences, epo
 
     for epoch_num in range(0, epochs):
         print('\n=== Epoch {:} / {:} ==='.format(epoch_num + 1, epochs))
+        
         model.train()
 
         total_acc_train = 0
@@ -182,16 +189,22 @@ def train(model, train_dataloader, dev_dict_dataloader, test_dict_sentences, epo
         
         # Dev results for each corpus. We don't need to save the results.
         for corpus in dev_dict_dataloader:
-            _ = get_predictions(model, corpus, dev_dict_dataloader[corpus])
+            _ = get_predictions(model, 
+                                corpus, 
+                                dev_dict_dataloader[corpus])
             
         # we want the results of specific epochs for specific corpora. 
         # we define the epochs and the corpora and we save only these results.
         
         if epoch_num+1 in specific_results:
-            if corpus in specific_results[epoch_num+1]:
-                test_results = get_predictions(model, corpus, dev_dict_dataloader[corpus], 
+            for corpus in specific_results[epoch_num+1]:
+                test_results = get_predictions(model, 
+                                               corpus, 
+                                               test_dict_dataloader[corpus], 
                                                print_results=False)
-                print_results_to_file(corpus, test_dict_sentences[corpus], test_results,
+                print_results_to_file(corpus, 
+                                      test_dict_sentences[corpus], 
+                                      test_results,
                                       inv_mappings, substitutions_file)
 
                 
@@ -200,13 +213,26 @@ def train(model, train_dataloader, dev_dict_dataloader, test_dict_sentences, epo
 print('\nModel: ', args.transformer_model)
 print('Batch size: ', args.batch_size * args.gradient_accumulation_steps)
 print('\nStart training...\n')
-train(model, train_dataloader, dev_dict_dataloader, test_dict_sentences, args.num_epochs, specific_results)
+train(model, 
+      train_dataloader,
+      dev_dict_dataloader, 
+      test_dict_sentences, 
+      test_dict_dataloader,
+      args.num_epochs, 
+      specific_results)
 print('\nTraining Done!')
 
 
 # ------- Testing ---------
 
+print('Testing...')
 for corpus in test_dict_dataloader:
-    test_results = get_predictions(model, corpus, test_dict_dataloader[corpus])
-    print_results_to_file(corpus, test_dict_sentences[corpus], test_results,
-                          inv_mappings, substitutions_file)
\ No newline at end of file
+    test_results = get_predictions(model, 
+                                   corpus, 
+                                   test_dict_dataloader[corpus]
+                                  )
+#     print_results_to_file(corpus, 
+#                           test_dict_sentences[corpus], 
+#                           test_results,
+#                           inv_mappings, 
+#                           substitutions_file)
\ No newline at end of file
diff --git a/utils.py b/utils.py
index 2e0635712fd1b91c88a0feff61cca2667f77eeb6..aa80eeabe4455324dd7e40ccf4045c97dc14258d 100644
--- a/utils.py
+++ b/utils.py
@@ -141,7 +141,10 @@ def open_sentences(path_to_corpora, mappings_dict):
 # Testing functions
 # ===============
 
-def get_predictions(model, corpus, test_dataloader, print_results=True):
+def get_predictions(model,
+                    corpus, 
+                    test_dataloader, 
+                    print_results=True):
     
     ''' Function to get the model's predictions for one corpus' test set.
         Can print accuracy using scikit-learn.
@@ -183,7 +186,10 @@ def get_predictions(model, corpus, test_dataloader, print_results=True):
     return all_preds
     
     
-def get_predictions_huggingface(trainer, corpus, test_set, print_results=True):
+def get_predictions_huggingface(trainer,
+                                corpus, 
+                                test_set, 
+                                print_results=True):
     
     ''' SPECIFI FUNCTION FOR THE HUGGINGFACE TRAINER.
         Function to get the model's predictions for one corpus' test set.
@@ -203,8 +209,11 @@ def get_predictions_huggingface(trainer, corpus, test_set, print_results=True):
     return preds
     
     
-def print_results_to_file(corpus, test_sentences, test_results, 
-                          inv_mappings_dict, substitutions_file):
+def print_results_to_file(corpus, 
+                          test_sentences, 
+                          test_results, 
+                          inv_mappings_dict, 
+                          substitutions_file):
     
     ''' Function to print a new file with the test predictions per 
         the specifications of the Shared task.