diff --git a/codes/models/audio/music/transformer_diffusion12.py b/codes/models/audio/music/transformer_diffusion12.py
index 0ca0c029..96c3f385 100644
--- a/codes/models/audio/music/transformer_diffusion12.py
+++ b/codes/models/audio/music/transformer_diffusion12.py
@@ -342,7 +342,8 @@ class TransformerDiffusionWithQuantizer(nn.Module):
         # higher gradients. Ideally we would use parameter groups, but ZeroRedundancyOptimizer makes this trickier than
         # directly fiddling with the gradients.
         for p in scaled_grad_parameters:
-            p.grad *= .2
+            if hasattr(p, 'grad') and p.grad is not None:
+                p.grad *= .2
 
 
 class TransformerDiffusionWithARPrior(nn.Module):