def METHOD_NAME(self) -> Mapping[str, Mapping[int, str]]:
    if self.task == 'multiple-choice':
        dynamic_axis = {0: 'batch', 1: 'choice', 2: 'sequence'}
    else:
        dynamic_axis = {0: 'batch', 1: 'sequence'}
    return OrderedDict([
        ('input_ids', dynamic_axis),
        ('attention_mask', dynamic_axis),
    ])