Merge pull request #768 from take-cheeze/disable_tf32_cudnn_onnx

[onnx] Disable cudnn tf32 in torchvision test
pfnet · Oct 24, 2023 · c5b4d58 · c5b4d58
2 parents 3f023d3 + abe1248
commit c5b4d58
Showing 1 changed file with 11 additions and 6 deletions.
diff --git 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