(self, x)
| 690 | self.embedding.weight.data.uniform_(-1.0 / n_e, 1.0 / n_e) |
| 691 | |
| 692 | def forward(self, x): |
| 693 | d = ( |
| 694 | torch.sum(x**2, 1, keepdim=True) |
| 695 | + torch.sum(self.embedding.weight**2, 1) |
| 696 | - 2 * torch.matmul(x, self.embedding.weight.T) |
| 697 | ) |
| 698 | min_indicies = torch.argmin(d, 1) |
| 699 | z_q = self.embedding(min_indicies) |
| 700 | return z_q, min_indicies |
| 701 | |
| 702 | |
| 703 | class Quantizer(torch.nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected