Commit 13e9dc3b authored by Phil Wang's avatar Phil Wang
Browse files

fix ability to load either model or ema model from trainer saved object, for... 

fix ability to load either model or ema model from trainer saved object, for @djqualia in https://github.com/lucidrains/audiolm-pytorch/issues/89
parent 535fa4a1

audiolm_pytorch/soundstream.py

+7 −0
Original line number Diff line number Diff line
@@ -498,6 +498,13 @@ class SoundStream(nn.Module): 
        assert path.exists()

        self.load_state_dict(torch.load(str(path)))



    def load_from_trainer_saved_obj(self, path, ema = False):

        key = 'ema_model' if ema else 'model'

        path = Path(path)

        assert path.exists()

        trainer_obj = torch.load(str(path))

        self.load_state_dict(trainer_obj[key])



    def non_discr_parameters(self):

        return [

            *self.encoder.parameters(),

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

  version = '0.12.4',

  license='MIT',

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

  author = 'Phil Wang',