Reduce flaky tests in test_vmap via skips and tolerances

Author: zou3519

test/test_vmap.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from typing import OrderedDict
-from unittest.case import skipIf, skip
+from unittest.case import skipIf
 from torch.testing._internal.common_utils import TestCase, run_tests
 import torch
 import torch.nn.functional as F
@@ -29,6 +29,7 @@
 from common_utils import (
     get_fallback_and_vmap_exhaustive,
     xfail,
+    skip,
     skipOps,
     check_vmap_fallback,
     tol1,
@@ -1066,7 +1067,7 @@ def func3(x, y, z, w):
 
         assert expected.allclose(out)
 
-    @skip("Somehow, vmap and autocast do not work on CPU")
+    @unittest.skip("Somehow, vmap and autocast do not work on CPU")
     def test_vmap_autocast_cpu(self):
         self._test_vmap_autocast("cpu")
 
@@ -3127,6 +3128,7 @@ class TestVmapOperatorsOpInfo(TestCase):
         xfail('column_stack', ''),
         xfail('pca_lowrank', ''),
         xfail('svd_lowrank', ''),
+        skip('linalg.eigh', ''),  # Flaky but is likely a real problem
 
         # required rank 4 tensor to use channels_last format
         xfail('bfloat16'),
@@ -3145,8 +3147,10 @@ class TestVmapOperatorsOpInfo(TestCase):
     @opsToleranceOverride('TestVmapOperatorsOpInfo', 'test_vmap_exhaustive', (
         tol1('linalg.det',
              {torch.float32: tol(atol=1e-04, rtol=1e-04)}, device_type='cuda'),
+        # The following is often flaky, but just on windows.
+        # We should investigate if it's actually a problem or not.
         tol1('nn.functional.conv_transpose3d',
-             {torch.float32: tol(atol=1.5e-04, rtol=1e-04)}, device_type='cuda'),
+             {torch.float32: tol(atol=1e-04, rtol=1e-02)}, device_type='cuda'),
     ))
     @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)})
     @skipOps('TestVmapOperatorsOpInfo', 'test_vmap_exhaustive', vmap_fail)