diff --git a/tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py b/tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py index 4c49ceb0..8c0a8cce 100644 --- a/tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py +++ b/tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py @@ -14,12 +14,17 @@ @pytest.mark.filterwarnings("ignore:Converting a tensor to a Python boolean might cause the trace to be incorrect:torch.jit.TracerWarning") def test_eval_resnet18(): - run_model_test( - torchvision.models.resnet.resnet18(**resnet18_kwargs), - (torch.rand(1, 3, 224, 224),), - rtol=1e-03, - use_gpu=True, - ) + old_allow_tf32 = torch.backends.cudnn.allow_tf32 + try: + torch.backends.cudnn.allow_tf32 = False + run_model_test( + torchvision.models.resnet.resnet18(**resnet18_kwargs), + (torch.rand(1, 3, 224, 224),), + rtol=1e-03, + use_gpu=True, + ) + finally: + torch.backends.cudnn.allow_tf32 = old_allow_tf32 @pytest.mark.gpu