Update hf_model.py

hf_model.py

@@ -182,7 +182,7 @@ class Triad(nn.Module):
         assert image is not None or audio is not None or text_list is not None, "At least one modality must be provided"
         if image is not None: assert image is not str, "Frames should be a path to an image"
         if audio is not None: 
-            assert isinstance(audio, torch.Tensor) and audio.shape[0] == 1 and len(audio.shape) == 2, "Audio must be a PyTorch tensor of shape (1, T)"
+            assert isinstance(audio, torch.Tensor) and len(audio.shape) == 2, "Audio must be a PyTorch tensor of shape (B, T)"
         if text_list is not None:
             assert isinstance(text_list, list) and len(text_list) == 1, "Text list must be a list of strings of length 1"
         if image is not None: