diff --git a/tests/test_rand_affine_grid.py b/tests/test_rand_affine_grid.py index 60ac40f468..722bafb0e5 100644 --- a/tests/test_rand_affine_grid.py +++ b/tests/test_rand_affine_grid.py @@ -18,7 +18,7 @@ from monai.transforms import RandAffineGrid from tests.utils import TEST_NDARRAYS, assert_allclose, is_tf32_env -_rtol = 1e-1 if is_tf32_env else 1e-4 +_rtol = 1e-1 if is_tf32_env() else 1e-4 TESTS = [] for p in TEST_NDARRAYS: diff --git a/tests/utils.py b/tests/utils.py index c4f9bd1b70..1a547fc2d2 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -586,13 +586,11 @@ def _wrapper(*args, **kwargs): def _cache_original_func(obj) -> None: """cache the original function by name, so that the decorator doesn't shadow it.""" - global _original_funcs _original_funcs[obj.__name__] = obj def _del_original_func(obj): """pop the original function from cache.""" - global _original_funcs _original_funcs.pop(obj.__name__, None) if torch.cuda.is_available(): # clean up the cached function torch.cuda.synchronize()