Function _check_device

tests/python/pytorch/dataloading/test_dataloader.py:215–223

(data)

def _check_device(data):
    if isinstance(data, dict):
        for k, v in data.items():
            assert v.device == F.ctx()
    elif isinstance(data, list):
        for v in data:
            assert v.device == F.ctx()
    else:
        assert data.device == F.ctx()


@pytest.mark.parametrize("sampler_name", ["full", "neighbor"])