def METHOD_NAME(batch_size, num_labels, num_channels=3, height=224, width=224):
    pixel_values = torch.randn(batch_size,
                               num_channels,
                               height,
                               width,
                               device=torch.cuda.current_device(),
                               dtype=torch.float)
    labels = torch.randint(0, num_labels, (batch_size,), device=torch.cuda.current_device(), dtype=torch.int64)
    return dict(pixel_values=pixel_values, labels=labels)