def get_src_tokens(self, correct_prediction, aggregate):
        """
            correct_prediction: True if the net_output (src_tokens) should
            predict the correct target
            aggregate: True if the criterion expects net_output (src_tokens)
            aggregated across time axis
        """
        predicted_idx = 0 if correct_prediction else 1
        if aggregate:
            src_tokens = torch.zeros((2, 2), dtype=torch.float)
            for b in range(2):
                src_tokens[b][predicted_idx] = 1.0
        else:
            src_tokens = torch.zeros((2, 10, 2), dtype=torch.float)
            for b in range(2):
                for t in range(10):
                    src_tokens[b][t][predicted_idx] = 1.0
        return src_tokens

            correct_prediction: True if the net_output (src_tokens) should
            predict the correct target
            aggregate: True if the criterion expects net_output (src_tokens)
            aggregated across time axis
        
infoxlm/fairseq/tests/speech_recognition/asr_test_base.py