| 219 | |
| 220 | # Necessary for my KFAC implementation. |
| 221 | class AddBias(nn.Module): |
| 222 | def __init__(self, bias): |
| 223 | super(AddBias, self).__init__() |
| 224 | self._bias = nn.Parameter(bias.unsqueeze(1)) |
| 225 | |
| 226 | def forward(self, x): |
| 227 | if x.dim() == 2: |
| 228 | bias = self._bias.t().view(1, -1) |
| 229 | else: |
| 230 | bias = self._bias.t().view(1, -1, 1, 1) |
| 231 | |
| 232 | return x + bias |
| 233 | |
| 234 | |
| 235 | def init(module, weight_init, bias_init, gain=1): |