diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 95261007..453e0503 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -921,12 +921,19 @@ If you want each process to load the full dataset, ignore this warning. if isinstance(batch, torch.Tensor): return batch.cuda(gpu_id) - # when list/tuple - elif isinstance(batch, list) or isinstance(batch, tuple): + # when list + elif isinstance(batch, list): for i, x in enumerate(batch): batch[i] = self.transfer_batch_to_gpu(x, gpu_id) return batch + # when tuple + elif isinstance(batch, tuple): + batch = list(batch) + for i, x in enumerate(batch): + batch[i] = self.transfer_batch_to_gpu(x, gpu_id) + return tuple(batch) + # when dict elif isinstance(batch, dict): for k, v in batch.items(): diff --git a/tests/test_models.py b/tests/test_models.py index 4813c542..59fa5a45 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -92,6 +92,18 @@ def test_single_gpu_batch_parse(): assert batch[0]['a'].device.index == 0 and batch[0]['a'].type() == 'torch.cuda.FloatTensor' assert batch[0]['b'].device.index == 0 and batch[0]['b'].type() == 'torch.cuda.FloatTensor' + # tuple of tensor list and list of tensor dict + batch = ([torch.rand(2, 3) for _ in range(2)], + [{'a': torch.rand(2, 3), 'b': torch.rand(2, 3)} for _ in range(2)]) + batch = trainer.transfer_batch_to_gpu(batch, 0) + assert batch[0][0].device.index == 0 and batch[0][0].type() == 'torch.cuda.FloatTensor' + + assert batch[1][0]['a'].device.index == 0 + assert batch[1][0]['a'].type() == 'torch.cuda.FloatTensor' + + assert batch[1][0]['b'].device.index == 0 + assert batch[1][0]['b'].type() == 'torch.cuda.FloatTensor' + def test_early_stopping_cpu_model(): """