diff --git a/Configuration/config.ini b/Configuration/config.ini
index d6c860553e48b27f2c16d5d1c2228c86e51b7fc4..d28938d20a106d1d2b13082329053a41caec7945 100644
--- a/Configuration/config.ini
+++ b/Configuration/config.ini
@@ -12,15 +12,14 @@ max_atoms_in_one_type=250
 dim_encoder = 768
 
 [MODEL_DECODER]
-dim_decoder = 16
-num_rnn_layers=1
+dim_decoder = 32
 dropout=0.1
 teacher_forcing=0.05
 
 [MODEL_LINKER]
 nhead=4
 dim_feedforward=246
-dim_embedding_atoms=16
+dim_embedding_atoms=32
 dim_polarity_transfo=128
 layer_norm_eps=1e-5
 dropout=0.1
@@ -29,6 +28,6 @@ sinkhorn_iters=3
 [MODEL_TRAINING]
 device=cpu
 batch_size=32
-epoch=20
+epoch=25
 seed_val=42
 learning_rate=0.005
\ No newline at end of file
diff --git a/Linker/Linker.py b/Linker/Linker.py
index c986c0e1bf3b96cf479a864d4e952a6cf47fa6af..b4a5c8014d9d03c1974d4034b6e7d908afc2129a 100644
--- a/Linker/Linker.py
+++ b/Linker/Linker.py
@@ -189,10 +189,10 @@ class Linker(Module):
             checkpoint_dir, writer = output_create_dir()
 
         for epoch_i in range(epochs):
-            epoch_acc, epoch_loss = self.train_epoch(training_dataloader)
+            avg_train_loss, avg_accuracy_train = self.train_epoch(training_dataloader)
 
-            print("Average Loss on train dataset : ", epoch_loss)
-            print("Average Accuracy on train dataset : ", epoch_acc)
+            print("Average Loss on train dataset : ", avg_train_loss)
+            print("Average Accuracy on train dataset : ", avg_accuracy_train)
 
             if checkpoint:
                 self.__checkpoint_save(
@@ -200,20 +200,20 @@ class Linker(Module):
 
             if validation_rate > 0.0:
                 with torch.no_grad():
-                    accuracy_test, average_test_loss = self.eval_epoch(validation_dataloader, self.cross_entropy_loss)
-                    print("Average Loss on test dataset : ", average_test_loss)
+                    loss_test, accuracy_test = self.eval_epoch(validation_dataloader, self.cross_entropy_loss)
+                    print("Average Loss on test dataset : ", loss_test)
                     print("Average Accuracy on test dataset : ", accuracy_test)
 
             if tensorboard:
                 writer.add_scalars(f'Accuracy', {
-                    'Train': epoch_acc}, epoch_i)
+                    'Train': avg_accuracy_train}, epoch_i)
                 writer.add_scalars(f'Loss', {
-                    'Train': epoch_loss}, epoch_i)
+                    'Train': avg_train_loss}, epoch_i)
                 if validation_rate > 0.0:
                     writer.add_scalars(f'Accuracy', {
                         'Validation': accuracy_test}, epoch_i)
                     writer.add_scalars(f'Loss', {
-                        'Validation': average_test_loss}, epoch_i)
+                        'Validation': loss_test}, epoch_i)
 
             print('\n')
 
@@ -336,7 +336,7 @@ class Linker(Module):
         accuracy = mesure_accuracy(batch_true_links, axiom_links_pred)
         loss = cross_entropy_loss(logits_axiom_links_pred, batch_true_links)
 
-        return accuracy, loss
+        return loss, accuracy
 
     def eval_epoch(self, dataloader, cross_entropy_loss):
         r"""Average the evaluation of all the batch.
@@ -347,11 +347,11 @@ class Linker(Module):
         accuracy_average = 0
         loss_average = 0
         for step, batch in enumerate(dataloader):
-            accuracy, loss = self.eval_batch(batch, cross_entropy_loss)
+            loss, accuracy = self.eval_batch(batch, cross_entropy_loss)
             accuracy_average += accuracy
-            loss_average += loss
+            loss_average += float(loss)
 
-        return accuracy_average / len(dataloader), loss_average / len(dataloader)
+        return loss_average / len(dataloader), accuracy_average / len(dataloader)
 
     def load_weights(self, model_file):
         print("#" * 15)
diff --git a/train.py b/train.py
index 571777ac750eef16c30044cddd4a203abf57e522..513e384b6cae6d1c66b44faad81ba71202186830 100644
--- a/train.py
+++ b/train.py
@@ -6,7 +6,7 @@ from utils import read_csv_pgbar
 
 torch.cuda.empty_cache()
 batch_size = int(Configuration.modelTrainingConfig['batch_size'])
-nb_sentences = batch_size * 200
+nb_sentences = batch_size * 20
 epochs = int(Configuration.modelTrainingConfig['epoch'])
 
 file_path_axiom_links = 'Datasets/goldANDsilver_dataset_links.csv'