Add dimension and shape check (#2563)

Summary:
Don't allow users to input incorrect dimensions

Pull Request resolved: https://github.com/pytorch/audio/pull/2563

Reviewed By: carolineechen

Differential Revision: D38074360

Pulled By: skim0514

fbshipit-source-id: 7bcae515706eb358ca6f68c50c7c0ccace1c3f95

torchaudio/prototype/models/hdemucs.py

@@ -516,8 +516,14 @@ class HDemucs(torch.nn.Module):
                 output tensor split into sources of shape `(batch_size, num_sources, channel, num_frames)`
         """
 
-        if len(input.shape) == 2:
-            input = input.view(1, self.audio_channels, -1)
+        if input.ndim != 3:
+            raise ValueError(f"Expected 3D tensor with dimensions (batch, channel, frames). Found: {input.shape}")
+
+        if input.shape[1] != self.audio_channels:
+            raise ValueError(
+                f"The channel dimension of input Tensor must match `audio_channels` of HDemucs model. "
+                f"Found:{input.shape[1]}."
+            )
 
         x = input
         length = x.shape[-1]