examples/research_projects/jax-projects/big_bird/bigbird_flax.py/DataCollator/_fetch_inputs
class DataCollator:    def _fetch_inputs(self, input_ids: list):
        attention_mask = [1 for _ in range(len(input_ids))]
        while len(input_ids) < self.max_length:
            input_ids.append(self.pad_id)
            attention_mask.append(0)
        return input_ids, attention_mask