From 37076fdb6acab99fc8bd2f718703ab4d6abc31b3 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Thu, 30 Nov 2023 08:35:07 -0800
Subject: [PATCH] address
 https://github.com/lucidrains/audiolm-pytorch/issues/256

---
 audiolm_pytorch/trainer.py | 1 +
 audiolm_pytorch/version.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/audiolm_pytorch/trainer.py b/audiolm_pytorch/trainer.py
index 66ace6c..00e3c23 100644
--- a/audiolm_pytorch/trainer.py
+++ b/audiolm_pytorch/trainer.py
@@ -536,6 +536,7 @@ def train_step(self):
                     accum_log(logs, {name: discr_loss.item() / self.grad_accum_every})
 
         if exists(self.discr_max_grad_norm):
+            self.accelerator.clip_grad_norm_(self.soundstream.discriminators.parameters(), self.discr_max_grad_norm)
             self.accelerator.clip_grad_norm_(self.soundstream.stft_discriminator.parameters(), self.discr_max_grad_norm)
 
         # gradient step for all discriminators
diff --git a/audiolm_pytorch/version.py b/audiolm_pytorch/version.py
index cfe6447..fa2822c 100644
--- a/audiolm_pytorch/version.py
+++ b/audiolm_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.8.3'
+__version__ = '1.8.4'