From 4ba7b0f6ccb658fec12316417aab1e52a33882ee Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Mon, 28 Dec 2020 16:45:38 +0100 Subject: [PATCH] Torch import guard. --- tests/test_pipelines_conversational.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/test_pipelines_conversational.py b/tests/test_pipelines_conversational.py index 17f8b6c0..88bb276d 100644 --- a/tests/test_pipelines_conversational.py +++ b/tests/test_pipelines_conversational.py @@ -14,13 +14,22 @@ import unittest -from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Conversation, ConversationalPipeline, pipeline -from transformers.models.gpt2 import GPT2Config, GPT2LMHeadModel +from transformers import ( + AutoModelForSeq2SeqLM, + AutoTokenizer, + Conversation, + ConversationalPipeline, + is_torch_available, + pipeline, +) from transformers.testing_utils import require_torch, slow, torch_device from .test_pipelines_common import DummyTok, MonoInputPipelineCommonMixin +if is_torch_available(): + from transformers.models.gpt2 import GPT2Config, GPT2LMHeadModel + DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0