From 2dd0b9e6e96e5a09b78721d5933e008d6f5cb9b4 Mon Sep 17 00:00:00 2001 From: James Betker Date: Sun, 22 May 2022 12:25:45 -0600 Subject: [PATCH] mel_head should be optional --- codes/models/audio/music/flat_diffusion.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/codes/models/audio/music/flat_diffusion.py b/codes/models/audio/music/flat_diffusion.py index e90196b7..0d6a343a 100644 --- a/codes/models/audio/music/flat_diffusion.py +++ b/codes/models/audio/music/flat_diffusion.py @@ -271,6 +271,8 @@ class FlatDiffusion(nn.Module): assert not (return_code_pred and precomputed_code_embeddings is not None), "I cannot compute a code_pred output for you." unused_params = [] + if not return_code_pred: + unused_params.extend(list(self.mel_head.parameters())) if conditioning_free: code_emb = self.unconditioned_embedding.repeat(x.shape[0], 1, x.shape[-1]) unused_params.extend(list(self.code_converter.parameters()) + list(self.code_embedding.parameters()))