From 45afefabedddb1fb266bff9652af22eac3bc1f7d Mon Sep 17 00:00:00 2001 From: James Betker Date: Sun, 24 Jul 2022 18:00:14 -0600 Subject: [PATCH] fix booboo --- codes/models/audio/music/transformer_diffusion13.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/models/audio/music/transformer_diffusion13.py b/codes/models/audio/music/transformer_diffusion13.py index 45f4383c..f2fc1080 100644 --- a/codes/models/audio/music/transformer_diffusion13.py +++ b/codes/models/audio/music/transformer_diffusion13.py @@ -273,7 +273,7 @@ class TransformerDiffusion(nn.Module): return groups def before_step(self, step): - scaled_grad_parameters = list(itertools.chain.from_iterable([lyr.out.parameters() for lyr in self.diff.layers])) + scaled_grad_parameters = list(itertools.chain.from_iterable([lyr.out.parameters() for lyr in self.layers])) # Scale back the gradients of the blkout and prenorm layers by a constant factor. These get two orders of magnitudes # higher gradients. Ideally we would use parameter groups, but ZeroRedundancyOptimizer makes this trickier than # directly fiddling with the gradients.