def test_device_consistency() -> None:
    """Test if the pro metric yields the same results between cpu and gpu.

    Note: This test will only run on a GPU-enabled device.
    """
    transform = RandomAffine(5, None, (0.95, 1.05), 5)

    batch = torch.zeros((32, 256, 256))
    for i in range(batch.shape[0]):
        batch[i, ...] = generate_perlin_noise(256, 256, scale=(4, 4)) > 0.5
    # ground truth mask is int type
    batch = batch.type(torch.int32)

    preds = transform(batch)

    pro_cpu = PRO()
    pro_gpu = PRO()

    pro_cpu.update(preds.cpu(), batch.cpu())
    pro_gpu.update(preds.cuda(), batch.cuda())

    assert torch.isclose(pro_cpu.compute(), pro_gpu.compute().cpu())
Test if the pro metric yields the same results between cpu and gpu.

    Note: This test will only run on a GPU-enabled device.
    
tests/unit/metrics/test_pro.py