    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
tests/speech_recognition/asr_test_base.py