import torch


def get_detection_input(batch_size=1, img_width=224, img_height=224):
    """
    Sample input for detection models, usable for tracing or testing
    """
    return (
            torch.rand(batch_size, 3, img_width, img_height),
            torch.arange(0, batch_size).long(),
            torch.Tensor([1, 1, 200, 200]).repeat((batch_size, 1)),
            torch.full((batch_size,), 1).long(),
    )