(self, input_dim: int, codebook_size: int, codebook_dim: int)
| 350 | """ |
| 351 | |
| 352 | def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int): |
| 353 | super().__init__() |
| 354 | self.codebook_size = codebook_size |
| 355 | self.codebook_dim = codebook_dim |
| 356 | |
| 357 | self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1) |
| 358 | self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1) |
| 359 | self.codebook = nn.Embedding(codebook_size, codebook_dim) |
| 360 | |
| 361 | def forward(self, z): |
| 362 | """Quantized the input tensor using a fixed codebook and returns |