x: (B,C,T,F) conv_cache: (B,C, (kT-1)*8, F) tra_cache: (3,1,B,C)
(self, x, en_outs, conv_cache, tra_cache)
| 277 | ]) |
| 278 | |
| 279 | def forward(self, x, en_outs, conv_cache, tra_cache): |
| 280 | """ |
| 281 | x: (B,C,T,F) |
| 282 | conv_cache: (B,C, (kT-1)*8, F) |
| 283 | tra_cache: (3,1,B,C) |
| 284 | """ |
| 285 | x, conv_cache[:,:, 6:16, :], tra_cache[0] = self.de_convs[0](x + en_outs[4], conv_cache[:,:, 6:16, :], tra_cache[0]) |
| 286 | x, conv_cache[:,:, 2:6, :], tra_cache[1] = self.de_convs[1](x + en_outs[3], conv_cache[:,:, 2:6, :], tra_cache[1]) |
| 287 | x, conv_cache[:,:, :2, :], tra_cache[2] = self.de_convs[2](x + en_outs[2], conv_cache[:,:, :2, :], tra_cache[2]) |
| 288 | |
| 289 | for i in range(3, 5): |
| 290 | x = self.de_convs[i](x + en_outs[4-i]) |
| 291 | return x, conv_cache, tra_cache |
| 292 | |
| 293 | |
| 294 | class Mask(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected