(self, samples, mask)
| 151 | # broadcast_tensors(self.buffers()) |
| 152 | |
| 153 | def replace_(self, samples, mask): |
| 154 | modified_codebook = torch.where( |
| 155 | mask[..., None], sample_vectors(samples, self.codebook_size), self.embed |
| 156 | ) |
| 157 | self.embed.data.copy_(modified_codebook) |
| 158 | |
| 159 | def expire_codes_(self, batch_samples): |
| 160 | if self.threshold_ema_dead_code == 0: |
no test coverage detected