From b317c68ac967d57fbac7155c95bdb757b37ce0a3 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 20 May 2022 11:12:53 -0600 Subject: [PATCH] fix --- codes/models/audio/music/flat_diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/models/audio/music/flat_diffusion.py b/codes/models/audio/music/flat_diffusion.py index f0bcaf08..ee9de790 100644 --- a/codes/models/audio/music/flat_diffusion.py +++ b/codes/models/audio/music/flat_diffusion.py @@ -183,7 +183,7 @@ class FlatDiffusion(nn.Module): groups = { 'minicoder': list(self.contextual_embedder.parameters()), 'layers': list(self.layers.parameters()), - 'code_converters': list(self.code_embedding.parameters()) + list(self.code_converter.parameters()) + list(self.latent_conditioner.parameters()) + list(self.latent_conditioner.parameters()), + 'code_converters': list(self.embeddings.parameters())) + list(self.latent_conditioner.parameters()), 'timestep_integrator': list(self.conditioning_timestep_integrator.parameters()) + list(self.integrating_conv.parameters()), 'time_embed': list(self.time_embed.parameters()), }