Repository: Lightning-AI/pytorch-lightning

Issue:
Training on GPU failed with Torchtext when using include_lengths=True in torchtext.data.Field
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
The issues raises in pytorch_lightning/utilities/apply_func.py which assumes that the attributes of a Batch from trochtext are Tensors, however if torchtext.data.Field is configured to include a length Tensor (include_lengths=True) the field is a tuple.

A bugfix is prepared and a PR can be submitted soon.

### To Reproduce

Steps to reproduce the behavior:

1. Use Torchtext Field with include_lengths=True on a GPU machine and fit model.
2. Training works on CPU but fails on GPU with: `TypeError: cannot unpack non-iterable NoneType object`

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

### Full Error Message
```
Traceback (most recent call last):
 File "debug_torchtext.py", line 105, in <module>
  trainer.fit(model)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1003, in fit
  results = self.single_gpu_train(model)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 186, in single_gpu_train
  results = self.run_pretrain_routine(model)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1213, in run_pretrain_routine
  self.train()
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 370, in train
  self.run_training_epoch()
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 452, in run_training_epoch
  batch_output = self.run_training_batch(batch, batch_idx)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 632, in run_training_batch
  self.hiddens
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 776, in optimizer_closure
  hiddens)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 944, in training_forward
  batch = self.transfer_batch_to_gpu(batch, gpu_id)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 159, in transfer_batch_to_gpu
  return self.__transfer_batch_to_device(batch, device)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 164, in __transfer_batch_to_device
  return model.transfer_batch_to_device(batch, device)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/core/hooks.py", line 242, in transfer_batch_to_device
  return move_data_to_device(batch, device)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 128, in move_data_to_device
  return apply_to_collection(batch, dtype=(TransferableDataType, Batch), function=batch_to)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 35, in apply_to_collection
  return function(data, *args, **kwargs)
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 103, in batch_to
  device_field = getattr(data, field).to(device, non_blocking=True)
AttributeError: 'tuple' object has no attribute 'to'
Exception ignored in: <function tqdm.__del__ at 0x7fcb5e0b2680>
Traceback (most recent call last):
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/tqdm/std.py", line 1086, in __del__
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/tqdm/std.py", line 1293, in close
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/tqdm/std.py", line 1471, in display
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/tqdm/std.py", line 1089, in __repr__
 File "/home1/thschaaf/miniconda3/envs/p37/lib/python3.7/site-packages/tqdm/std.py", line 1433, in format_dict
TypeError: cannot unpack non-iterable NoneType object
```

#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->
```
import torch
from torch import nn, Tensor
import pytorch_lightning as pl
from pytorch_lightning import Trainer, seed_everything
from torchtext import data
seed_everything(1234)
def get_debug_data_loader():
    text_field = data.Field(sequential=True, pad_first=False,
                            init_token="
", eos_token="
", include_lengths=True)
    example1 = data.example.Example.fromdict({"text": "a b c a c"}, {"text": ("text", text_field)})
    example2 = data.example.Example.fromdict({"text": "b c a a"}, {"text": ("text", text_field)})
    example3 = data.example.Example.fromdict({"text": "c b a"}, {"text": ("text", text_field)})
    dataset = data.Dataset([example1, example2, example3], {"text": text_field})
    text_field.build_vocab(dataset)
    iterator = data.Iterator(dataset, batch_size=3,
                             sort_key=None, device=None, batch_size_fn=None,
                             train=True, repeat=False, shuffle=None, sort=None, sort_within_batch=None)
    return iterator, text_field
class DebugModel(pl.LightningModule):
    def __init__(self):
        super(DebugModel, self).__init__()
        # setup data loader
        self.debug_data_loader, self.text_field = get_debug_data_loader()
        self.learning_rate = 0.001
        self.hid_dim = 4
        pad_idx = self.text_field.vocab.stoi['<pad>']
        self.criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)
        self.INPUT_DIM = len(self.text_field.vocab)
        self.ENC_EMB_DIM = 4  # keep it small for debugging
        self.embedding = nn.Embedding(self.INPUT_DIM, self.ENC_EMB_DIM)
        self.rnn = nn.GRU(self.ENC_EMB_DIM, self.hid_dim, 1, bidirectional=False)
        self.out = nn.Linear(self.hid_dim, self.embedding.num_embeddings)
        self.OUTPUT_DIM = len(self.text_field.vocab)
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
    def forward(self, input_seq, length):
        embedded: Tensor = self.embedding(input_seq)
        packed_embedded: Tensor = torch.nn.utils.rnn.pack_padded_sequence(embedded, length, batch_first=False,
                                                                          enforce_sorted=False)
        packed_outputs, hidden = self.rnn(packed_embedded)  # [sent len, batch size, emb dim]
        outputs, length = torch.nn.utils.rnn.pad_packed_sequence(packed_outputs)
        # outputs -> [sent len, batch size, hid dim * n directions]
        # hidden -> [n layers * n directions, batch size, hid dim]
        output = outputs.squeeze(0)
        prediction = self.out(output)
        return prediction
    @staticmethod
    def _parse_batch(batch):
        source = batch.text[0]
        source_length = batch.text[1]
        return source, source_length
    def training_step(self, batch, batch_nb):
        x = self._parse_batch(batch)
        target, target_length = x
        output = self.forward(target, target_length)
        loss = self.criterion(output[:-1].view(-1, output.shape[2]), target[1:].view(-1))
        prefix = 'train'
        tensorboard_logs = {f'{prefix}_loss': loss.item()}
        result = {'loss': loss, 'log': tensorboard_logs}
        return result
    def train_dataloader(self):
        return self.debug_data_loader
model = DebugModel()
cuda_device_cnt = torch.cuda.device_count()
if cuda_device_cnt > 0:
    use_num_cuda_devices = 1
else:
    use_num_cuda_devices = None
trainer = Trainer(fast_dev_run=False, max_steps=None,
                  gradient_clip_val=10,
                  weights_summary='full', gpus=use_num_cuda_devices,
                  show_progress_bar=True)
trainer.fit(model)
```

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
Should not raise an error :-) 

### Environment

```
 CUDA:
    - GPU:
        - TITAN X (Pascal)
    - available:     True
    - version:      10.2
* Packages:
    - numpy:       1.17.3
    - pyTorch_debug:   False
    - pyTorch_version:  1.5.1
    - pytorch-lightning: 0.8.5
    - tensorboard:    2.2.2
    - tqdm:       4.47.0
* System:
    - OS:        Linux
    - architecture:
        - 64bit
        - 
    - processor:     x86_64
    - python:      3.7.4
    - version:      #1 SMP Tue Mar 17 23:49:17 UTC 2020
```
### Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
--- a/pytorch_lightning/utilities/apply_func.py
+++ b/pytorch_lightning/utilities/apply_func.py
@@ -6,6 +6,7 @@
 import torch
 
 import importlib
+
 TORCHTEXT_AVAILABLE = importlib.util.find_spec("torchtext") is not None
 if TORCHTEXT_AVAILABLE:
     from torchtext.data import Batch
@@ -92,6 +93,7 @@ def move_data_to_device(batch: Any, device: torch.device):
         - :meth:`torch.Tensor.to`
         - :class:`torch.device`
     """
+
     def batch_to(data):
         # try to move torchtext data first
         if TORCHTEXT_AVAILABLE and isinstance(data, Batch):
@@ -99,11 +101,10 @@ def batch_to(data):
             # Shallow copy because each Batch has a reference to Dataset which contains all examples
             device_data = copy(data)
             for field in data.fields:
-                # Batch contains output of Field.process(...) which is tensor hence .to(...) exists
-                device_field = getattr(data, field).to(device, non_blocking=True)
+                device_field = move_data_to_device(getattr(data, field), device)
                 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)


Test patch:
diff --git a/tests/base/datamodules.py b/tests/base/datamodules.py
--- a/tests/base/datamodules.py
+++ b/tests/base/datamodules.py
@@ -13,7 +13,7 @@ def __init__(self, data_dir: str = './'):
     def prepare_data(self):
         TrialMNIST(self.data_dir, train=True, download=True)
         TrialMNIST(self.data_dir, train=False, download=True)
-    
+
     def setup(self):
         mnist_full = TrialMNIST(root=self.data_dir, train=True, num_samples=64, download=True)
         self.mnist_train, self.mnist_val = random_split(mnist_full, [128, 64])
diff --git a/tests/callbacks/test_model_checkpoint.py b/tests/callbacks/test_model_checkpoint.py
--- a/tests/callbacks/test_model_checkpoint.py
+++ b/tests/callbacks/test_model_checkpoint.py
@@ -28,7 +28,7 @@ def test_model_checkpoint_with_non_string_input(tmpdir, save_top_k):
         max_epochs=2,
     )
     trainer.fit(model)
-    assert checkpoint.dirpath == tmpdir / trainer.logger.name / f'version_0' / 'checkpoints'
+    assert checkpoint.dirpath == tmpdir / trainer.logger.name / 'version_0' / 'checkpoints'
 
 
 @pytest.mark.parametrize(
diff --git a/tests/utilities/test_apply_func_torchtext.py b/tests/utilities/test_apply_func_torchtext.py
new file mode 100644
--- /dev/null
+++ b/tests/utilities/test_apply_func_torchtext.py
@@ -0,0 +1,52 @@
+import pytest
+import torch
+import torchtext
+from torchtext.data.example import Example
+
+from pytorch_lightning.utilities.apply_func import move_data_to_device
+
+
+def _get_torchtext_data_iterator(include_lengths=False):
+    text_field = torchtext.data.Field(sequential=True, pad_first=False,  # nosec
+                                      init_token="
", eos_token="
",  # nosec
+                                      include_lengths=include_lengths)  # nosec
+
+    example1 = Example.fromdict({"text": "a b c a c"}, {"text": ("text", text_field)})
+    example2 = Example.fromdict({"text": "b c a a"}, {"text": ("text", text_field)})
+    example3 = Example.fromdict({"text": "c b a"}, {"text": ("text", text_field)})
+
+    dataset = torchtext.data.Dataset(
+        [example1, example2, example3],
+        {"text": text_field},
+    )
+    text_field.build_vocab(dataset)
+
+    iterator = torchtext.data.Iterator(dataset, batch_size=3,
+                                       sort_key=None, device=None,
+                                       batch_size_fn=None,
+                                       train=True, repeat=False, shuffle=None,
+                                       sort=None, sort_within_batch=None)
+    return iterator, text_field
+
+
+@pytest.mark.parametrize('include_lengths', [False, True])
+@pytest.mark.parametrize(['device'], [pytest.param(torch.device('cuda', 0))])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test assumes GPU machine")
+def test_batch_move_data_to_device_torchtext_include_lengths(include_lengths, device):
+    data_iterator, _ = _get_torchtext_data_iterator(include_lengths=include_lengths)
+    data_iter = iter(data_iterator)
+    batch = next(data_iter)
+    batch_on_device = move_data_to_device(batch, device)
+
+    if include_lengths:
+        # tensor with data
+        assert (batch_on_device.text[0].device == device)
+        # tensor with length of data
+        assert (batch_on_device.text[1].device == device)
+    else:
+        assert (batch_on_device.text.device == device)
+
+
+@pytest.mark.parametrize('include_lengths', [False, True])
+def test_batch_move_data_to_device_torchtext_include_lengths_cpu(include_lengths):
+    test_batch_move_data_to_device_torchtext_include_lengths(include_lengths, torch.device('cpu'))