(
self,
d_model: int = 64,
strides: list = [2, 4, 8, 8],
d_latent: int = 64,
)
| 719 | |
| 720 | class Encoder(nn.Module): |
| 721 | def __init__( |
| 722 | self, |
| 723 | d_model: int = 64, |
| 724 | strides: list = [2, 4, 8, 8], |
| 725 | d_latent: int = 64, |
| 726 | ): |
| 727 | super().__init__() |
| 728 | # Create first convolution |
| 729 | self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)] |
| 730 | |
| 731 | # Create EncoderBlocks that double channels as they downsample by `stride` |
| 732 | for stride in strides: |
| 733 | d_model *= 2 |
| 734 | self.block += [EncoderBlock(d_model, stride=stride)] |
| 735 | |
| 736 | # Create last convolution |
| 737 | self.block += [ |
| 738 | Snake1d(d_model), |
| 739 | WNConv1d(d_model, d_latent, kernel_size=3, padding=1), |
| 740 | ] |
| 741 | |
| 742 | # Wrap black into nn.Sequential |
| 743 | self.block = nn.Sequential(*self.block) |
| 744 | self.enc_dim = d_model |
| 745 | |
| 746 | def forward(self, x): |
| 747 | return self.block(x) |
nothing calls this directly
no test coverage detected