import torch


def get_image(width, height, grayscale=False):
    """Generate image Tensor, returns CHW"""
    channels = 1 if grayscale else 3
    numel = width * height * channels
    img = torch.arange(numel, dtype=torch.int64) % 256
    img = img.reshape(channels, height, width).to(torch.uint8)
    return img