Avoid required torchtext import

Lightning-AI · Nov 30, 2021 · 6b814fa · 6b814fa
1 parent c13b159
commit 6b814fa
Showing 1 changed file with 12 additions and 4 deletions.
diff --git 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)