| 134 | |
| 135 | |
| 136 | class Decoder(nn.Module): |
| 137 | def __init__(self, layers, norm_layer=None, projection=None): |
| 138 | super(Decoder, self).__init__() |
| 139 | self.layers = nn.ModuleList(layers) |
| 140 | self.norm = norm_layer |
| 141 | self.projection = projection |
| 142 | |
| 143 | def forward(self, x, cross, x_mask=None, cross_mask=None): |
| 144 | for layer in self.layers: |
| 145 | x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask) |
| 146 | |
| 147 | if self.norm is not None: |
| 148 | x = self.norm(x) |
| 149 | |
| 150 | if self.projection is not None: |
| 151 | x = self.projection(x) |
| 152 | return x |