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),
        ]
    )