(module, weight_init, bias_init, gain=1)
| 10 | |
| 11 | |
| 12 | def init(module, weight_init, bias_init, gain=1): |
| 13 | weight_init(module.weight.data, gain=gain) |
| 14 | if module.bias is not None: |
| 15 | bias_init(module.bias.data) |
| 16 | else: |
| 17 | print('None bias') |
| 18 | return module |
| 19 | |
| 20 | |
| 21 | class Flatten(nn.Module): |