needs to be zero centered GP

lucidrains · Jan 12, 2025 · 2154a74 · 2154a74
1 parent 6740714
commit 2154a74
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/audiolm_pytorch/soundstream.py b/audiolm_pytorch/soundstream.py
@@ -67,7 +67,7 @@ def hinge_gen_loss(fake):
 def leaky_relu(p = 0.1):
     return nn.LeakyReLU(p)
 
-def gradient_penalty(wave, output, weight = 10):
+def gradient_penalty(wave, output, weight = 10, center = 0.):
     batch_size, device = wave.shape[0], wave.device
 
     gradients = torch_grad(
@@ -80,7 +80,7 @@ def gradient_penalty(wave, output, weight = 10):
     )[0]
 
     gradients = rearrange(gradients, 'b ... -> b (...)')
-    return weight * ((vector_norm(gradients, dim = 1) - 1) ** 2).mean()
+    return weight * ((vector_norm(gradients, dim = 1) - center) ** 2).mean()
 
 # better sequential
 

diff --git a/audiolm_pytorch/version.py b/audiolm_pytorch/version.py
@@ -1 +1 @@
-__version__ = '2.3.0'
+__version__ = '2.3.1'