Init method based on N(0, sigma/sqrt(2*num_layers).
(sigma, num_layers)
| 89 | return init_ |
| 90 | |
| 91 | def scaled_init_method(sigma, num_layers): |
| 92 | """Init method based on N(0, sigma/sqrt(2*num_layers).""" |
| 93 | std = sigma / math.sqrt(2.0 * num_layers) |
| 94 | def init_(tensor, **kwargs): |
| 95 | return torch.nn.init.normal_(tensor, mean=0.0, std=std) |
| 96 | |
| 97 | return init_ |
| 98 | |
| 99 | @torch.jit.script |
| 100 | def gelu_impl(x): |