(self, keys, queue_ptr, queue)
| 89 | |
| 90 | @torch.no_grad() |
| 91 | def _dequeue_and_enqueue(self, keys, queue_ptr, queue): |
| 92 | # gather keys before updating queue |
| 93 | if self._use_ddp_or_ddp2(self.trainer): |
| 94 | keys = concat_all_gather(keys) |
| 95 | |
| 96 | batch_size = keys.shape[0] |
| 97 | |
| 98 | ptr = int(queue_ptr) |
| 99 | assert self.hparams.num_negatives % batch_size == 0 # for simplicity |
| 100 | |
| 101 | # replace the keys at ptr (dequeue and enqueue) |
| 102 | queue[:, ptr : ptr + batch_size] = keys.T |
| 103 | ptr = (ptr + batch_size) % self.hparams.num_negatives # move pointer |
| 104 | |
| 105 | queue_ptr[0] = ptr |
| 106 | |
| 107 | @torch.no_grad() |
| 108 | def _batch_shuffle_ddp(self, x): # pragma: no cover |
no test coverage detected