def _dequeue_and_enqueue()

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