From bd98f01aac5a9e1b1db55cb5744c1ffa69ff39e1 Mon Sep 17 00:00:00 2001
From: Alexandre Defossez <alexandre.defossez@gmail.com>
Date: Thu, 29 Sep 2022 16:40:53 +0200
Subject: [PATCH] Fixing bug #374, issue in pad1d

---
 demucs/__init__.py |  2 +-
 demucs/hdemucs.py  | 12 ++++++++++--
 docs/release.md    |  6 ++++++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/demucs/__init__.py b/demucs/__init__.py
index 810fa9bb..90a2d633 100644
--- a/demucs/__init__.py
+++ b/demucs/__init__.py
@@ -4,4 +4,4 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-__version__ = "3.0.6a1"
+__version__ = "3.0.6a2"
diff --git a/demucs/hdemucs.py b/demucs/hdemucs.py
index 39e4a578..f003cb79 100644
--- a/demucs/hdemucs.py
+++ b/demucs/hdemucs.py
@@ -23,13 +23,21 @@
 def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
     """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
     If this is the case, we insert extra 0 padding to the right before the reflection happen."""
+    x0 = x
     length = x.shape[-1]
     padding_left, padding_right = paddings
     if mode == 'reflect':
         max_pad = max(padding_left, padding_right)
         if length <= max_pad:
-            x = F.pad(x, (0, max_pad - length + 1))
-    return F.pad(x, paddings, mode, value)
+            extra_pad = max_pad - length + 1
+            extra_pad_right = min(padding_right, extra_pad)
+            extra_pad_left = extra_pad - extra_pad_right
+            paddings = (padding_left - extra_pad_left, padding_right - extra_pad_right)
+            x = F.pad(x, (extra_pad_left, extra_pad_right))
+    out = F.pad(x, paddings, mode, value)
+    assert out.shape[-1] == length + padding_left + padding_right
+    assert (out[..., padding_left: padding_left + length] == x0).all()
+    return out
 
 
 class ScaledEmbedding(nn.Module):
diff --git a/docs/release.md b/docs/release.md
index 7893933d..6b9ffd1f 100644
--- a/docs/release.md
+++ b/docs/release.md
@@ -1,5 +1,11 @@
 # Release notes for Demucs
 
+## V3.0.6a, TBC
+
+Option to customize output path of stems (@CarlGao4)
+
+Fixed bug in pad1d leading to failure sometimes.
+
 ## V3.0.5, 17th of August 2022
 
 Added `--segment` flag to customize the segment length and use less memory (thanks @CarlGao4).