From 536c8558ae50c55d639aabd52f0e8c2d4cd94cef Mon Sep 17 00:00:00 2001 From: James Betker Date: Sat, 28 May 2022 22:32:38 -0600 Subject: [PATCH] fix --- codes/models/audio/music/transformer_diffusion5.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/models/audio/music/transformer_diffusion5.py b/codes/models/audio/music/transformer_diffusion5.py index 2ae84345..106913e6 100644 --- a/codes/models/audio/music/transformer_diffusion5.py +++ b/codes/models/audio/music/transformer_diffusion5.py @@ -136,7 +136,7 @@ class TransformerDiffusion(nn.Module): groups = { 'contextual_embedder': list(self.conditioning_embedder.parameters()), 'layers': list(self.layers.parameters()) + list(self.inp_block.parameters()), - 'code_converters': list(self.embeddings.parameters()) + list(self.code_converter.parameters()), + 'code_converters': list(self.input_converter.parameters()) + list(self.code_converter.parameters()), 'time_embed': list(self.time_embed.parameters()), } return groups