From ca3b0f461d69ec9899fc9100887d3454b8d1f4cf Mon Sep 17 00:00:00 2001 From: Matthias Fey Date: Sun, 26 May 2024 10:01:36 +0200 Subject: [PATCH] Fix `SentenceTransformer` test (#9362) --- test/nn/nlp/test_sentence_transformer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/nn/nlp/test_sentence_transformer.py b/test/nn/nlp/test_sentence_transformer.py index 24caa0e59fe2..d778d5abc1ae 100644 --- 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)