diff --git a/tests/helpers/imports.py b/tests/helpers/imports.py index fc86476f3f6ae..462b592502782 100644 --- a/tests/helpers/imports.py +++ b/tests/helpers/imports.py @@ -1,6 +1,14 @@ -from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8 +from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8, _TORCHTEXT_AVAILABLE -if _TORCH_GREATER_EQUAL_1_8: - from torchtext.legacy.data import Batch, Dataset, Example, Field, Iterator, LabelField +if _TORCHTEXT_AVAILABLE: + if _TORCH_GREATER_EQUAL_1_8: + from torchtext.legacy.data import Batch, Dataset, Example, Field, Iterator, LabelField + else: + from torchtext.data import Batch, Dataset, Example, Field, Iterator, LabelField else: - from torchtext.data import Batch, Dataset, Example, Field, Iterator, LabelField # noqa: F401 + Batch = type(None) + Dataset = type(None) + Example = type(None) + Field = type(None) + Iterator = type(None) + LabelField = type(None)