Update modeling_custom.py

modeling_custom.py

@@ -85,10 +85,11 @@ class CustomOutput(ModelOutput):
 
 class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
     def __init__(self, config):
+        config.torch_dtype = torch.bfloat16
         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, trust_remote_code=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)
         self.num_labels = config.num_labels
         config_dict = config.to_dict()
         self.num_objectives = config_dict.get("num_objectives", 19)