Fix padding (#50)

EleutherAI · Aug 14, 2024 · 51b011e · 51b011e
1 parent 352de4a
commit 51b011e
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/amt/data.py b/amt/data.py
@@ -111,12 +111,9 @@ def get_wav_segments(
             yield buffer
 
     if pad_last == True:
-        yield torch.cat(
-            (
-                buffer[stride_samples:],
-                torch.zeros(stride_samples, dtype=torch.float32),
-            ),
-            dim=0,
+        yield torch.nn.functional.pad(
+            buffer[stride_samples:],
+            (0.0, chunk_samples - len(buffer[stride_samples:])),
         )