(self, x)
| 173 | return x |
| 174 | |
| 175 | def quantize(self, x): |
| 176 | embed = self.embed.t() |
| 177 | dist = -( |
| 178 | x.pow(2).sum(1, keepdim=True) |
| 179 | - 2 * x @ embed |
| 180 | + embed.pow(2).sum(0, keepdim=True) |
| 181 | ) |
| 182 | embed_ind = dist.max(dim=-1).indices |
| 183 | return embed_ind |
| 184 | |
| 185 | def postprocess_emb(self, embed_ind, shape): |
| 186 | return embed_ind.view(*shape[:-1]) |