Made the warnings filter more specific

2025-09-03 11:16:51 -07:00
parent 84eb046a38
commit 2fcdc24699
1 changed files with 4 additions and 1 deletions
--- a/ucsinfer/main.py
+++ b/ucsinfer/main.py
@@ -246,7 +246,10 @@ def evaluate(dataset, offset, limit, model, no_foley):
 if __name__ == '__main__':
    os.environ['TOKENIZERS_PARALLELISM'] = 'false'

+    # sentence_transformers generates an error in PyTorch upon loading
    import warnings
-    warnings.simplefilter(action='ignore', category=FutureWarning)
+    warnings.filterwarnings(action='ignore', module='torch', 
+                            category=FutureWarning, 
+                            message=r"`encoder_attention_mask` is deprecated.*")

    ucsinfer()