Function test_pytorch_containers

dali/test/python/test_external_source_impl_utils.py:83–96

()

@attr("pytorch")
def test_pytorch_containers():
    import torch

    samples_cpu = [
        (torch.tensor(test_array), test_array),
    ]
    batches_cpu = [
        ([torch.tensor(test_array)], [test_array]),
        ([torch.tensor(test_array)] * 4, [test_array] * 4),
    ]
    disallowed_samples = [
        torch.tensor(test_array).cuda(),
    ]
    yield from run_checks(samples_cpu, batches_cpu, disallowed_samples, [])


@attr("cupy")