diff --git a/codes/models/audio/music/transformer_diffusion5.py b/codes/models/audio/music/transformer_diffusion5.py
index fc938086..e9ef9a23 100644
--- a/codes/models/audio/music/transformer_diffusion5.py
+++ b/codes/models/audio/music/transformer_diffusion5.py
@@ -218,6 +218,7 @@ class TransformerDiffusionWithQuantizer(nn.Module):
         del self.m2v.m2v.projector.projection
         del self.m2v.project_hid
         del self.m2v.project_q
+        del self.m2v.m2v.masked_spec_embed
 
     def update_for_step(self, step, *args):
         self.internal_step = step