diff --git a/docs/source/prototype.models.rst b/docs/source/prototype.models.rst new file mode 100644 index 0000000000..489f82e817 --- /dev/null +++ b/docs/source/prototype.models.rst @@ -0,0 +1,11 @@ +torchaudio.prototype.models +=========================== + +.. py:module:: torchaudio.prototype.models +.. currentmodule:: torchaudio.prototype.models + +conformer_rnnt_base +~~~~~~~~~~~~~~~~~~~ + +.. autofunction:: conformer_rnnt_base + diff --git a/docs/source/prototype.rst b/docs/source/prototype.rst index a38550f5af..677e37c4a8 100644 --- a/docs/source/prototype.rst +++ b/docs/source/prototype.rst @@ -19,4 +19,5 @@ imported explicitly, e.g. .. toctree:: prototype.io prototype.ctc_decoder + prototype.models prototype.pipelines diff --git a/torchaudio/prototype/models/rnnt.py b/torchaudio/prototype/models/rnnt.py index b53335f3af..ec28e82de1 100644 --- a/torchaudio/prototype/models/rnnt.py +++ b/torchaudio/prototype/models/rnnt.py @@ -41,6 +41,12 @@ def infer( def conformer_rnnt_base(): + r"""Builds basic version of Conformer RNN-T model. + + Returns: + RNNT: + Conformer RNN-T model. + """ encoder = _ConformerTranscriber() decoder = _Predictor( num_symbols=1024,