nicolinho
/

QRM-Llama3.1-8B

Model card Files Files and versions Community

nicolinho commited on Sep 27

Commit

74280ee

•

1 Parent(s): 710e02a

Update modeling_custom.py

Files changed (1) hide show

modeling_custom.py +1 -1

modeling_custom.py CHANGED Viewed

@@ -89,7 +89,7 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
         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 = 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)

         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 = 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)