Commit 5058cc44 authored by Phil Wang's avatar Phil Wang
Browse files

address two bugs, thanks to @apoorv2904 in issue... 

address two bugs, thanks to @apoorv2904  in issue https://github.com/lucidrains/audiolm-pytorch/issues/48
parent 4c3c3534

audiolm_pytorch/soundstream.py

+6 −1
Original line number Diff line number Diff line
@@ -355,7 +355,12 @@ class SoundStream(nn.Module): 
        self.load_state_dict(torch.load(str(path)))



    def non_discr_parameters(self):

        return [*self.encoder.parameters(), *self.decoder.parameters()]

        return [

            *self.encoder.parameters(),

            *self.decoder.parameters(),

            *self.encoder_attn.parameters(),

            *self.decoder_attn.parameters()

        ]



    @property

    def seq_len_multiple_of(self):

audiolm_pytorch/trainer.py

+2 −1
Original line number Diff line number Diff line
@@ -319,6 +319,8 @@ class SoundStreamTrainer(nn.Module): 


        # update discriminator



        self.discr_optim.zero_grad()



        for _ in range(self.grad_accum_every):

            wave, = next(self.dl_iter)

            wave = wave.to(device)
@@ -340,7 +342,6 @@ class SoundStreamTrainer(nn.Module): 
        # gradient step for all discriminators



        self.discr_optim.step()

        self.discr_optim.zero_grad()



        for ind in range(len(self.soundstream.discriminators)):

            discr_optimizer = getattr(self, f'multiscale_discr_optimizer_{ind}')

setup.py

+1 −1
Original line number Diff line number Diff line
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages 
setup(

  name = 'audiolm-pytorch',

  packages = find_packages(exclude=[]),

  version = '0.3.5',

  version = '0.3.6',

  license='MIT',

  description = 'AudioLM - Language Modeling Approach to Audio Generation from Google Research - Pytorch',

  author = 'Phil Wang',