def forward(self, model, sample, reduce=True):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        model.eval()
        all_retrieval_result = {}
        random.seed(0)
        for context_length in range(self.cfg.interval, self.cfg.tokens_per_sample + 1, self.cfg.interval):
            all_length = (context_length - 150)
            local_retrieval_result = []
            depth_number = 10
            for depth_ratio in range(0, depth_number + 1):
                prefix_length = int(all_length * depth_ratio / depth_number)
                suffix_length = all_length - prefix_length
                n_correct = 0
                times = 10
                for _ in range(times):
                    prompt, pass_key = self.generate_prompt_landmark(prefix_length, suffix_length)
                    prompt_tokens = self.task.tokenizer.encode(prompt, bos=True)
                    prompt_tokens = torch.tensor([prompt_tokens], device="cuda")
                    print(prompt_tokens.shape)
                    output = self.generate(model, prompt_tokens)
                    pred = self.task.tokenizer.decode(output[0, prompt_tokens.shape[1]:])
                    print("Answer: ", pass_key)
                    print("Pred: ", pred)
                    if pass_key in pred:
                        n_correct += 1
                local_retrieval_result.append(n_correct / times)
            all_retrieval_result[context_length] = local_retrieval_result

        print(all_retrieval_result)
        return 0, 1, {"loss": 0}
Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        
YOCO/yoco/criterions/needle_haystack.py