def METHOD_NAME(self) -> Optional[Dict[str, NeuralType]]:
    return {
        'loss': NeuralType(None, FloatType(), optional=True),
        'hidden_states': NeuralType(('B', 'T', 'D'), ChannelType()),
    }