Commit e4d31ce4 authored by Phil Wang's avatar Phil Wang
Browse files

actually save optim state for soundstream

parent 89a35002

audiolm_pytorch/trainer.py

+2 −0
Original line number Diff line number Diff line
@@ -229,6 +229,7 @@ class SoundStreamTrainer(nn.Module): 
        pkg = dict(

            model = self.soundstream.state_dict(),

            ema_model = self.ema_soundstream.state_dict(),

            optim = self.optim.state_dict(),

            discr_optim = self.discr_optim.state_dict()

        )
@@ -245,6 +246,7 @@ class SoundStreamTrainer(nn.Module): 


        self.soundstream.load_state_dict(pkg['model'])

        self.ema_soundstream.load_state_dict(pkg['ema_model'])

        self.optim.load_state_dict(pkg['optim'])

        self.discr_optim.load_state_dict(pkg['discr_optim'])



        for key, _ in self.multiscale_discriminator_iter():

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.2.1',

  version = '0.2.2',

  license='MIT',

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

  author = 'Phil Wang',