Function get_random_torch_batch

dali/test/python/test_dltensor_operator.py:300–309

(g, shapes, dtype)

def get_random_torch_batch(g, shapes, dtype):
    is_fp = torch.is_floating_point(torch.tensor([], dtype=dtype))
    if is_fp:
        return [torch.rand((shape), generator=g, dtype=dtype) for shape in shapes]
    else:
        iinfo = torch.iinfo(dtype)
        dtype_min, dtype_max = iinfo.min, iinfo.max
        return [
            torch.randint(dtype_min, dtype_max, shape, generator=g, dtype=dtype) for shape in shapes
        ]


def get_sliced_torch_case(case_name):