diff --git a/run_ner.py b/run_ner.py
index 2013432..9dec12e 100644
--- a/run_ner.py
+++ b/run_ner.py
@@ -667,7 +667,7 @@ def post_processing_function(examples, features, predictions, stage=f"eval"):
         eval_examples=eval_examples if training_args.do_eval else None,
         tokenizer=tokenizer,
         data_collator=data_collator,
-        callbacks=[EarlyStoppingCallback(early_stopping_patience=20)],
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=20)] if training_args.do_eval else None,
         post_process_function=post_processing_function,
         compute_metrics=None,
     )
diff --git a/src/config.py b/src/config.py
index 1157f15..38a67ea 100644
--- a/src/config.py
+++ b/src/config.py
@@ -1,5 +1,6 @@
 from transformers import PretrainedConfig
 
+
 class BinderConfig(PretrainedConfig):
 
     def __init__(
@@ -19,21 +20,11 @@ def __init__(
         threshold_loss_weight=0.5,
         ner_loss_weight=0.5,
     ):
-        self.pretrained_model_name_or_path=pretrained_model_name_or_path
-        self.cache_dir=cache_dir
-        self.revision=revision
-        self.use_auth_token=use_auth_token
-        self.hidden_dropout_prob=hidden_dropout_prob
-        self.max_span_width = max_span_width
-        self.use_span_width_embedding = use_span_width_embedding
-        self.linear_size = linear_size
-        self.init_temperature = init_temperature
-        self.start_loss_weight = start_loss_weight
-        self.end_loss_weight = end_loss_weight
-        self.span_loss_weight = span_loss_weight
-        self.threshold_loss_weight = threshold_loss_weight
-        self.ner_loss_weight = ner_loss_weight
-
+        self.pretrained_model_name_or_path = pretrained_model_name_or_path
+        self.cache_dir = cache_dir
+        self.revision = revision
+        self.use_auth_token = use_auth_token
+        self.hidden_dropout_prob = hidden_dropout_prob
         self.max_span_width = max_span_width
         self.use_span_width_embedding = use_span_width_embedding
         self.linear_size = linear_size
diff --git a/src/trainer.py b/src/trainer.py
index e19383c..f3bc3ae 100644
--- a/src/trainer.py
+++ b/src/trainer.py
@@ -33,9 +33,9 @@ def __post_init__(self):
             self.type_token_type_ids = torch.tensor(self.type_token_type_ids)
 
     def __call__(self, features: List) -> Dict[str, Any]:
-        batch = {}
-        batch['input_ids'] = torch.tensor([f['input_ids'] for f in features], dtype=torch.long)
-        batch['attention_mask'] = torch.tensor([f['attention_mask'] for f in features], dtype=torch.bool)
+        batch = {'input_ids': torch.tensor([f['input_ids'] for f in features], dtype=torch.long),
+                 'attention_mask': torch.tensor([f['attention_mask'] for f in features], dtype=torch.bool)}
+
         if "token_type_ids" in features[0]:
             batch['token_type_ids'] = torch.tensor([f['token_type_ids'] for f in features], dtype=torch.long)
 
@@ -48,7 +48,6 @@ def __call__(self, features: List) -> Dict[str, Any]:
             # For training
             ner = {}
             # Collate negative mask with shape [batch_size, num_types, ...].
-            start_negative_mask, end_negative_mask, span_negative_mask = [], [], []
             # [batch_size, num_types, seq_length]
             start_negative_mask = torch.tensor([f["ner"]["start_negative_mask"] for f in features], dtype=torch.bool)
             end_negative_mask = torch.tensor([f["ner"]["end_negative_mask"] for f in features], dtype=torch.bool)
@@ -59,7 +58,7 @@ def __call__(self, features: List) -> Dict[str, Any]:
             end_negative_mask[:, :, 0] = 1
             span_negative_mask[:, :, 0, 0] = 1
 
-            ner['start_negative_mask'] =  start_negative_mask
+            ner['start_negative_mask'] = start_negative_mask
             ner['end_negative_mask'] = end_negative_mask
             ner['span_negative_mask'] = span_negative_mask
 
@@ -152,7 +151,6 @@ def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metr
 
         return metrics
 
-
     def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_key_prefix: str = "test"):
         predict_dataloader = self.get_test_dataloader(predict_dataset)