import torch


def gen_images(size=128, batch_size=32, depth=4, device='cuda'):
    """Generate byte-type NHWC images."""
    tensor_size = (batch_size, size, size, depth * 3)
    # must make it float type so we can use uniform_()
    data_float = torch.empty(tensor_size, dtype=torch.float, device=device)
    data_float.uniform_(0, 256)
    data_byte = data_float.byte()
    return data_byte