Update modeling_custom.py

modeling_custom.py

@@ -87,8 +87,9 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
     def __init__(self, config):
         config.torch_dtype = torch.bfloat16
         super().__init__(config)
-        self.model = AutoModelForSequenceClassification.from_pretrained(
-                "Skywork/Skywork-Reward-Llama-3.1-8B", num_labels=1, torch_dtype=torch.bfloat16, use_flash_attention_2=True).model
+        #self.model = AutoModelForSequenceClassification.from_pretrained(
+        #        "Skywork/Skywork-Reward-Llama-3.1-8B", num_labels=1, torch_dtype=torch.bfloat16, use_flash_attention_2=True).model
+        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)