quick fix for not unwrapping model before getting config

lucidrains · lucidrains · commit ff6086839a90 · 2023-03-09T07:55:51.000-08:00
diff --git a/audiolm_pytorch/trainer.py b/audiolm_pytorch/trainer.py
@@ -257,7 +257,7 @@ def save(self, path):
         pkg = dict(
             model = self.accelerator.get_state_dict(self.soundstream),
             optim = self.optim.state_dict(),
-            config = self.soundstream._configs,
+            config = self.unwrapped_soundstream._configs,
             discr_optim = self.discr_optim.state_dict(),
             version = __version__
         )
diff --git a/audiolm_pytorch/version.py b/audiolm_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.23.2'
+__version__ = '0.23.3'

Original file line number	Diff line number	Diff line change
`@@ -257,7 +257,7 @@ def save(self, path):`
`257`	`257`	`pkg = dict(`
`258`	`258`	`model = self.accelerator.get_state_dict(self.soundstream),`
`259`	`259`	`optim = self.optim.state_dict(),`
`260`		`- config = self.soundstream._configs,`
	`260`	`+ config = self.unwrapped_soundstream._configs,`
`261`	`261`	`discr_optim = self.discr_optim.state_dict(),`
`262`	`262`	`version = __version__`
`263`	`263`	`)`