Update spleeter.py

spleeter.py

@@ -7,18 +7,9 @@ from torch.nn import functional as F
 from tqdm import tqdm
 
 def batchify(tensor: Tensor, T: int) -> Tensor:
-    """
-    partition tensor into segments of length T, zero pad any ragged samples
-    Args:
-        tensor(Tensor): BxCxFxL
-    Returns:
-        tensor of size (B*[L/T] x C x F x T)
-    """
-    # Zero pad the original tensor to an even multiple of T
     orig_size = tensor.size(-1)
     new_size = math.ceil(orig_size / T) * T
     tensor = F.pad(tensor, [0, new_size - orig_size])
-    # Partition the tensor into multiple samples of length T and stack them into a batch
     return torch.cat(torch.split(tensor, T, dim=-1), dim=0)
 
 
@@ -223,7 +214,7 @@ class Splitter(nn.Module):
 
     def infer_with_batches(self, stft_mag, batch_size):
         masks = {name: [] for name in self.stems.keys()}
-        with torch.no_grad():
+        with torch.inference_mode():
             for i in tqdm(range(0, stft_mag.shape[0], batch_size)):
                 batch = stft_mag[i:i + batch_size]
                 batch_outputs = {name: net(batch) for name, net in self.stems.items()}
@@ -231,3 +222,4 @@ class Splitter(nn.Module):
                     masks[name].append(batch_outputs[name])
         masks = {name: torch.cat(masks[name], dim=0) for name in masks}
         return masks
+