def METHOD_NAME(self) -> Mapping[str, Mapping[int, str]]:
    if self.task == "image-classification":
        return OrderedDict([("logits", {0: "batch"})])
    else:
        return OrderedDict([("last_hidden_state", {0: "batch"}), ("pooler_output", {0: "batch"})])