Skip chunk if too quiet

remfx/models.py

@@ -84,6 +84,9 @@ class RemFX(pl.LightningModule):
         x, y, _, _ = batch  # x, y = (B, C, T), (B, C, T)
 
         loss, output = self.model((x, y))
+        # Crop target to match output
+        if output.shape[-1] < y.shape[-1]:
+            y = causal_crop(y, output.shape[-1])
         self.log(f"{mode}_loss", loss)
         # Metric logging
         with torch.no_grad():

remfx/utils.py

@@ -158,6 +158,9 @@ def select_random_chunk(
     max_len = audio.shape[-1] - new_chunk_size
     random_start = torch.randint(0, max_len, (1,)).item()
     chunk = audio[:, random_start : random_start + new_chunk_size]
+    # Skip if energy too low
+    if torch.mean(torch.abs(chunk)) < 1e-6:
+        return None
     resampled_chunk = torchaudio.functional.resample(chunk, sr, sample_rate)
     return resampled_chunk