diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py index b5ec664b2abd0..2003bca6fd2e1 100644 --- a/pytorch_lightning/utilities/apply_func.py +++ b/pytorch_lightning/utilities/apply_func.py @@ -1,10 +1,14 @@ from abc import ABC from collections import Mapping, Sequence +from copy import copy from typing import Any, Callable, Union import torch -from torchtext.data import Batch -from copy import copy + +import importlib +TORCHTEXT_AVAILABLE = importlib.util.find_spec("torchtext") is not None +if TORCHTEXT_AVAILABLE: + from torchtext.data import Batch def apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable, *args, **kwargs) -> Any: @@ -86,9 +90,10 @@ def move_data_to_device(batch: Any, device: torch.device): - :meth:`torch.Tensor.to` - :class:`torch.device` """ - def batch_to(data): - if isinstance(data, Batch): + # try to move torchtext data first + if TORCHTEXT_AVAILABLE and isinstance(data, Batch): + # Shallow copy because each Batch has a reference to Dataset which contains all examples device_data = copy(data) for field in data.fields: @@ -96,6 +101,7 @@ def batch_to(data): device_field = getattr(data, field).to(device, non_blocking=True) setattr(device_data, field, device_field) return device_data + else: + return data.to(device, non_blocking=True) - return data.to(device, non_blocking=True) return apply_to_collection(batch, dtype=(TransferableDataType, Batch), function=batch_to) diff --git a/requirements/base.txt b/requirements/base.txt index e045fab912919..64375637e1711 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -6,5 +6,4 @@ torch>=1.3 tensorboard>=1.14 future>=0.17.1 # required for builtins in setup.py # pyyaml>=3.13 -PyYAML>=5.1 # OmegaConf requirement -torchtext>=0.3.1 \ No newline at end of file +PyYAML>=5.1 # OmegaConf requirement \ No newline at end of file diff --git a/requirements/extra.txt b/requirements/extra.txt index a9d8b6bdf486d..71a9ea9ccc6d3 100644 --- a/requirements/extra.txt +++ b/requirements/extra.txt @@ -10,4 +10,5 @@ matplotlib>=3.1.1 horovod>=0.19.1 omegaconf>=2.0.0 # scipy>=0.13.3 -scikit-learn>=0.20.0 \ No newline at end of file +scikit-learn>=0.20.0 +torchtext>=0.3.1 \ No newline at end of file