in covidprognosis/models/moco_model.py [0:0]
def _dequeue_and_enqueue(self, keys: Tensor):
# gather keys before updating queue
keys = concat_all_gather(keys)
batch_size = keys.shape[0]
assert isinstance(self.queue_ptr, Tensor)
ptr = int(self.queue_ptr)
assert (
self.K % batch_size == 0
), f"batch_size={batch_size}, K={self.K}" # for simplicity
# replace the keys at ptr (dequeue and enqueue)
self.queue[:, ptr : ptr + batch_size] = keys.T
ptr = (ptr + batch_size) % self.K # move pointer
self.queue_ptr[0] = ptr