mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-10 12:21:57 +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:
@@ -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