| 91 | return h |
| 92 | |
| 93 | def reverse(self, output): |
| 94 | if self.training and self.initialized.item() == 0: |
| 95 | if not self.allow_reverse_init: |
| 96 | raise RuntimeError( |
| 97 | "Initializing ActNorm in reverse direction is " |
| 98 | "disabled by default. Use allow_reverse_init=True to enable." |
| 99 | ) |
| 100 | else: |
| 101 | self.initialize(output) |
| 102 | self.initialized.fill_(1) |
| 103 | |
| 104 | if len(output.shape) == 2: |
| 105 | output = output[:, :, None, None] |
| 106 | squeeze = True |
| 107 | else: |
| 108 | squeeze = False |
| 109 | |
| 110 | h = output / self.scale - self.loc |
| 111 | |
| 112 | if squeeze: |
| 113 | h = h.squeeze(-1).squeeze(-1) |
| 114 | return h |