mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-30 20:55:17 +00:00
079bf3cb26
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
15 lines
373 B
Python
15 lines
373 B
Python
import math
|
|
|
|
import torch
|
|
|
|
|
|
def init_normal(tensor, sigma):
|
|
"""Init method based on N(0, sigma)."""
|
|
torch.nn.init.normal_(tensor, mean=0.0, std=sigma)
|
|
|
|
|
|
def output_init_normal(tensor, sigma, num_layers):
|
|
"""Init method based on N(0, sigma/sqrt(2*num_layers)."""
|
|
std = sigma / math.sqrt(2.0 * num_layers)
|
|
torch.nn.init.normal_(tensor, mean=0.0, std=std)
|