In the forward pass of moco, i am measuring the shuffling time, that is
    def forward(self, im_q, im_k):
        (....)
        # compute key features
        with torch.no_grad():  # no gradient to keys
            self._momentum_update_key_encoder()  # update the key encoder
            
            start_time = time.time()
            # shuffle for making use of BN
            im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
            torch.cuda.synchronize()
            end_time = time.time()
            shuffle_time = end_time - start_time
            (...)