eval bug fix

2022-06-10 13:51:06 -06:00 · 2022-06-10 13:51:06 -06:00 · 89bd40d39f
commit 89bd40d39f
parent 84469f3538
1 changed files with 2 additions and 0 deletions
--- a/codes/models/audio/music/transformer_diffusion8.py
+++ b/codes/models/audio/music/transformer_diffusion8.py
@ -234,6 +234,8 @@ class TransformerDiffusionWithQuantizer(nn.Module):
        diff = self.diff(x, timesteps, codes=proj, conditioning_input=conditioning_input,
                         conditioning_free=conditioning_free)

+        if disable_diversity:
+            return diff
        if mse is None:
            return diff, diversity_loss
        return diff, diversity_loss, mse