From 5f84ee6370de316220b2ad697e36de78d9de0207 Mon Sep 17 00:00:00 2001 From: CohleM Date: Mon, 3 Feb 2025 15:50:50 +0545 Subject: [PATCH] fix weight initialization for residual block --- model/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/model.py b/model/model.py index 20535fb..883efc8 100644 --- a/model/model.py +++ b/model/model.py @@ -333,7 +333,7 @@ class Transformer(PreTrainedModel): self.apply(self._init_weights) for pn, p in self.named_parameters(): - if pn.endswith('w3.weight') or pn.endswith('wo.weight'): + if pn.endswith('w2.weight') or pn.endswith('wo.weight'): torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * params.n_layers)) self.last_loss = None