fix chunk size too small error

test_unet.py

@@ -92,8 +92,13 @@ def separate(
     with torch.no_grad():
         mag = X.abs()
         frames = 2584
+        min_frames = 516
         masked_tf_rep = []
-        for chunk in tqdm(mag.split(frames, dim=-1)):
+        chunks = mag.split(frames, dim=-1)
+        if len(chunks) > 1 and chunks[-1].shape[-1] < min_frames:
+            chunks[-2:] = torch.cat(chunks[-2:], -1).chunk(2, -1)
+
+        for chunk in tqdm(chunks):
             masked_tf_rep += [x_umx_target(chunk.unsqueeze(0)).squeeze()]
         masked_tf_rep = torch.cat(masked_tf_rep, -1) * mag