Function test_assert

tests/transforms/test_create_grid_and_affine.py:160–167

(func, params, expected)

def test_assert(func, params, expected):
    gpu_test = ("torch_gpu",) if torch.cuda.is_available() else ()
    for b in ("torch", "numpy") + gpu_test:
        if b == "torch_gpu":
            m = func(*params, device="cuda:0", backend="torch")
        else:
            m = func(*params, backend=b)
        assert_allclose(m, expected, type_test=False, rtol=1e-2 if is_tf32_env() else 1e-5, atol=1e-5)


class TestCreateAffine(unittest.TestCase):