Don't convert namedtuple to tuple (#1589)

* Don't convert namedtuple to tuple

* Test namedtuples sent to device correctly
This commit is contained in:
Nathan Breitsch
2020-04-30 08:04:50 -04:00
committed by GitHub
parent d40425d257
commit 3eac6cfd4f
2 changed files with 17 additions and 4 deletions
+8
View File
@@ -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."""