(self, x)
| 302 | return quantize |
| 303 | |
| 304 | def forward(self, x): |
| 305 | device = x.device |
| 306 | x = rearrange(x, "b d n -> b n d") |
| 307 | x = self.project_in(x) |
| 308 | |
| 309 | quantize, embed_ind = self._codebook(x) |
| 310 | |
| 311 | if self.training: |
| 312 | quantize = x + (quantize - x).detach() |
| 313 | |
| 314 | loss = torch.tensor([0.0], device=device, requires_grad=self.training) |
| 315 | |
| 316 | if self.training: |
| 317 | if self.commitment_weight > 0: |
| 318 | commit_loss = F.mse_loss(quantize.detach(), x) |
| 319 | loss = loss + commit_loss * self.commitment_weight |
| 320 | |
| 321 | quantize = self.project_out(quantize) |
| 322 | quantize = rearrange(quantize, "b n d -> b d n") |
| 323 | return quantize, embed_ind, loss |
| 324 | |
| 325 | |
| 326 | class ResidualVectorQuantization(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected