mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
Don't convert namedtuple to tuple (#1589)
* Don't convert namedtuple to tuple * Test namedtuples sent to device correctly
This commit is contained in:
@@ -461,10 +461,15 @@ class TrainerDPMixin(ABC):
|
||||
|
||||
# when tuple
|
||||
if isinstance(batch, tuple):
|
||||
batch = list(batch)
|
||||
for i, x in enumerate(batch):
|
||||
batch[i] = self.__transfer_data_to_device(x, device, gpu_id)
|
||||
return tuple(batch)
|
||||
# when namedtuple
|
||||
if hasattr(batch, '_fields'):
|
||||
elem_type = type(batch)
|
||||
return elem_type(*(self.__transfer_data_to_device(x, device, gpu_id) for x in batch))
|
||||
else:
|
||||
batch = list(batch)
|
||||
for i, x in enumerate(batch):
|
||||
batch[i] = self.__transfer_data_to_device(x, device, gpu_id)
|
||||
return tuple(batch)
|
||||
|
||||
# when dict
|
||||
if isinstance(batch, dict):
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from collections import namedtuple
|
||||
import platform
|
||||
|
||||
import pytest
|
||||
@@ -221,6 +222,13 @@ def test_single_gpu_batch_parse():
|
||||
assert batch[1][0]['b'].device.index == 0
|
||||
assert batch[1][0]['b'].type() == 'torch.cuda.FloatTensor'
|
||||
|
||||
# namedtuple of tensor
|
||||
BatchType = namedtuple('BatchType', ['a', 'b'])
|
||||
batch = [BatchType(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].a.device.index == 0
|
||||
assert batch[0].a.type() == 'torch.cuda.FloatTensor'
|
||||
|
||||
|
||||
def test_simple_cpu(tmpdir):
|
||||
"""Verify continue training session on CPU."""
|
||||
|
||||
Reference in New Issue
Block a user