implement fix and test (#3459)

Lightning-AI · Sep 11, 2020 · bd5f53c · bd5f53c
1 parent 93cf6d0
commit bd5f53c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
@@ -121,4 +121,5 @@ def batch_to(data):
         kwargs = dict(non_blocking=True) if isinstance(data, torch.Tensor) else {}
         return data.to(device, **kwargs)
 
-    return apply_to_collection(batch, dtype=(TransferableDataType, Batch), function=batch_to)
+    dtype = (TransferableDataType, Batch) if TORCHTEXT_AVAILABLE else TransferableDataType
+    return apply_to_collection(batch, dtype=dtype, function=batch_to)
diff --git a/tests/models/test_gpu.py b/tests/models/test_gpu.py
@@ -341,6 +341,12 @@ def test_single_gpu_batch_parse():
     trainer = Trainer(gpus=1)
     trainer.accelerator_backend = GPUBackend(trainer)
 
+    # non-transferrable types
+    primitive_objects = [None, {}, [], 1.0, "x", [None, 2], {"x": (1, 2), "y": None}]
+    for batch in primitive_objects:
+        data = trainer.accelerator_backend.batch_to_device(batch, torch.device('cuda:0'))
+        assert data == batch
+
     # batch is just a tensor
     batch = torch.rand(2, 3)
     batch = trainer.accelerator_backend.batch_to_device(batch, torch.device('cuda:0'))