NormHead-forward nn.Parameter

modeling_baichuan.py

@@ -504,7 +504,7 @@ class NormHead(nn.Module):
             norm_weight = nn.functional.normalize(self.weight)
         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