Update modeling_custom.py

modeling_custom.py

@@ -86,9 +86,9 @@ class CustomOutput(ModelOutput):
 class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        #self.model = AutoModelForSequenceClassification.from_pretrained(
-        #        "sfairXC/FsfairX-LLaMA3-RM-v0.1", num_labels=1, torch_dtype=torch.bfloat16, use_flash_attention_2=True)
-        self.model = LlamaModel(config).to(torch.bfloat16)
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+                "sfairXC/FsfairX-LLaMA3-RM-v0.1", num_labels=1, torch_dtype=torch.bfloat16, use_flash_attention_2=True, trust_remote_code=True,)
+        #self.model = LlamaModel(config).to(torch.bfloat16)
         self.num_labels = config.num_labels
         config_dict = config.to_dict()
         self.num_objectives = config_dict.get("num_objectives", 19)