Commit 2e301432 authored by Leon Wu's avatar Leon Wu
Browse files

Fix type hint mistake

parent 6e0dab4e

audiolm_pytorch/audiolm_pytorch.py

+2 −2
Original line number Diff line number Diff line
@@ -1268,7 +1268,7 @@ class CoarseTransformerWrapper(nn.Module): 
        self,

        *,

        transformer: CoarseTransformer,

        codec: Optional[SoundStream, EncodecWrapper]  = None,

        codec: Optional[Union[SoundStream, EncodecWrapper]]  = None,

        wav2vec: Optional[Union[FairseqVQWav2Vec, HubertWithKmeans]] = None,

        audio_conditioner: Optional[AudioConditionerBase] = None,

        pad_id = -1,
@@ -1484,7 +1484,7 @@ class FineTransformerWrapper(nn.Module): 
        self,

        *,

        transformer: FineTransformer,

        codec: Optional[SoundStream, EncodecWrapper] = None,

        codec: Optional[Union[SoundStream, EncodecWrapper]] = None,

        audio_conditioner: Optional[AudioConditionerBase] = None,

        coarse_cross_entropy_loss_weight = 1.,

        pad_id = -1,