fix NormHead eval

I encountered this problem when using Baichuan2-7B-Base with deepspeed stage3 for sft. A similar situation also happened in the place such as https://github.com/baichuan-inc/Baichuan2/issues/39#issuecomment-1710146497
I found that Baichuan2-13B-Chat has solved this problem, so I synced the code here

modeling_baichuan.py

@@ -502,9 +502,10 @@ class NormHead(nn.Module):
     def forward(self, hidden_states):
         if self.training:
             norm_weight = nn.functional.normalize(self.weight)
+            self.first_flag = False
         elif self.first_flag:
             self.first_flag = False
-            self.weight = nn.Parameter(nn.functional.normalize(self.weight))
+            self.weight.data = nn.functional.normalize(self.weight)
             norm_weight = self.weight
         else:
             norm_weight = self.weight