(layer)
| 39 | return x * (1 + scale.unsqueeze(1).unsqueeze(1)) + shift.unsqueeze(1).unsqueeze(1) |
| 40 | |
| 41 | def zero_init(layer): |
| 42 | nn.init.zeros_(layer.weight) |
| 43 | if layer.bias is not None: |
| 44 | nn.init.zeros_(layer.bias) |
| 45 | return layer |
| 46 | |
| 47 | |
| 48 | def checkpoint(function, *args, **kwargs): |