def METHOD_NAME(distributed):
    device = idist.device()
    def get_test_cases():
        y_pred = torch.randint(0, 256, (n_iters * batch_size, 3, 16, 16), device=device, dtype=torch.uint8)
        y = (y_pred * 0.65).to(torch.uint8)
        return y_pred, y
    n_iters = 100
    batch_size = 10
    rank = idist.get_rank()
    for i in range(3):
        # check multiple random inputs as random exact occurencies are rare
        torch.manual_seed(42 + rank + i)
        y_pred, y = get_test_cases()
        _test(y_pred, y, 100, "cpu", n_iters, batch_size, atol=1e-8)
        if device.type != "xla":
            _test(y_pred, y, 100, idist.device(), n_iters, batch_size, atol=1e-8)