diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index a638c0d1..0d710124 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -18,6 +18,7 @@
 
 
 """
+from typing import List, Dict
 from dataclasses import dataclass
 from typing import Optional, Union
 
@@ -35,7 +36,7 @@ class RankGenCollator:
     max_length: Optional[int] = None
     max_examples: Optional[int] = None
 
-    def __call__(self, batch: list[dict[str, str]]) -> dict[str, torch.Tensor]:
+    def __call__(self, batch: List[Dict[str, str]]) -> Dict[str, torch.Tensor]:
         prefixes = []
         better_answers = []
         worse_answers = []
diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py
index 9312f7d3..7500901b 100644
--- a/model/reward/instructor/trainer.py
+++ b/model/reward/instructor/trainer.py
@@ -137,7 +137,6 @@ if __name__ == "__main__":
         output_dir=f"{model_name}-finetuned",
         num_train_epochs=training_conf["num_train_epochs"],
         warmup_steps=500,
-        loss_function=training_conf["loss"],
         learning_rate=training_conf["learning_rate"],
         # half_precision_backend="apex",
         fp16=training_conf["fp16"],