def METHOD_NAME(self) -> Optional[Dict[str, NeuralType]]:
    return {
        "input_ids": NeuralType(('B', 'T'), ChannelType()),
        "token_type_ids": NeuralType(('B', 'T'), ChannelType(), optional=True),
        "attention_mask": NeuralType(('B', 'T'), MaskType(), optional=True),
        "labels": NeuralType(('B', 'T'), ChannelType(), optional=True),
        'past_key_values': [[NeuralType(None, StringType(), optional=True)]],
        'use_cache': NeuralType(None, VoidType(), optional=True),
        'position_ids': NeuralType(('B', 'T'), ChannelType(), optional=True),
        "return_dict": NeuralType(None, StringType(), optional=True),
        "output_attentions": NeuralType(None, StringType(), optional=True),
        "output_hidden_states": NeuralType(None, StringType(), optional=True),
        "max_length": NeuralType(None, IntType(), optional=True),
    }