Fix SentenceTransformer test (#9362)

pyg-team · May 26, 2024 · ca3b0f4 · ca3b0f4
1 parent d3ce05c
commit ca3b0f4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/nn/nlp/test_sentence_transformer.py b/test/nn/nlp/test_sentence_transformer.py
@@ -1,11 +1,12 @@
 import pytest
 
 from torch_geometric.nn.nlp import SentenceTransformer
-from torch_geometric.testing import onlyFullTest, withCUDA
+from torch_geometric.testing import onlyFullTest, withCUDA, withPackage
 
 
 @withCUDA
 @onlyFullTest
+@withPackage('transformers')
 @pytest.mark.parametrize('batch_size', [None, 1])
 def test_sentence_transformer(batch_size, device):
     model = SentenceTransformer(model_name='prajjwal1/bert-tiny').to(device)