c-my/transformers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Torch import guard.

Browse Source
This commit is contained in:
Nicolas Patry
2020-12-28 16:49:46 +01:00
parent f9a057392f
commit 4ba7b0f6cc
1 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_pipelines_conversational.py
+11 -2
View File
@@ -14,13 +14,22 @@
import unittest import unittest
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Conversation, ConversationalPipeline, pipeline from transformers import (
from transformers.models.gpt2 import GPT2Config, GPT2LMHeadModel AutoModelForSeq2SeqLM,
AutoTokenizer,
Conversation,
ConversationalPipeline,
is_torch_available,
pipeline,
)
from transformers.testing_utils import require_torch, slow, torch_device from transformers.testing_utils import require_torch, slow, torch_device
from .test_pipelines_common import DummyTok, MonoInputPipelineCommonMixin 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 DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0