Update inference.py

inference.py

@@ -26,6 +26,17 @@ import warnings
 warnings.filterwarnings("ignore")
 
 
+def disable_rnn_flatten_parameters(model: nn.Module) -> None:
+    """
+    Disable cuDNN RNN weight flattening to avoid inference-mode in-place updates.
+    This is a workaround for RuntimeError about inplace updates to inference tensors
+    when running RNNs during inference.
+    """
+    for module in model.modules():
+        if isinstance(module, (nn.RNN, nn.GRU, nn.LSTM)):
+            module.flatten_parameters = lambda *args, **kwargs: None
+
+
 def run_folder(model, args, config, device, verbose: bool = False):
     """
     Process a folder of audio files for source separation.
@@ -174,12 +185,18 @@ def proc_folder(dict_args):
         checkpoint = torch.load(args.start_check_point, weights_only=False, map_location='cpu')
         load_start_checkpoint(args, model, checkpoint, type_='inference')
 
+    # Workaround for cuDNN RNN flattening with inference tensors
+    disable_rnn_flatten_parameters(model)
+
     print("Instruments: {}".format(config.training.instruments))
 
     # in case multiple CUDA GPUs are used and --device_ids arg is passed
     if isinstance(args.device_ids, list) and len(args.device_ids) > 1 and not args.force_cpu:
         model = nn.DataParallel(model, device_ids=args.device_ids)
 
+    # Ensure flattened parameters are disabled on the wrapped model too
+    disable_rnn_flatten_parameters(model)
+
     model = model.to(device)
 
     print("Model load time: {:.2f} sec".format(time.time() - model_load_start_time))