(self, embed_id)
| 54 | ) # (codebook_size, D'), ensure fp32 |
| 55 | |
| 56 | def decode_code(self, embed_id): # embed_id: (B, T) |
| 57 | embed = ( |
| 58 | F.embedding(embed_id, self.codebook).transpose(1, 2).float() |
| 59 | ) # (B, D', T), ensure fp32 |
| 60 | return embed |
| 61 | |
| 62 | def encode_code(self, z: torch.Tensor): # z: (B, D, T) |
| 63 | # logging.info(f"{self.cluster_size = }, {self.codebook = }, {self.embed_avg = }, {self.inited = }") |
no outgoing calls
no test coverage detected