[release/2.9] Only skip linalg.eig assertion in test_torch_return_types_returns (#3095)

Author: ethanwee1

test/functorch/test_vmap.py

@@ -58,7 +58,6 @@
     onlyCUDA,
     OpDTypes,
     ops,
-    skipCUDAIfNoMagma,
     tol,
     toleranceOverride,
 )
@@ -71,6 +70,7 @@
     run_tests,
     skipIfTorchDynamo,
     subtest,
+    TEST_WITH_ROCM,
     TEST_WITH_TORCHDYNAMO,
     TestCase,
     unMarkDynamoStrictTest,
@@ -5043,7 +5043,6 @@ def op(inp, running_mean, running_var, weight, bias, training):
 
         test(self, op, tuple(inputs), in_dims=tuple(in_dims))
 
-    @skipCUDAIfNoMagma
     def test_torch_return_types_returns(self, device):
         t = torch.randn(3, 2, 2, device=device)
         self.assertTrue(
@@ -5057,9 +5056,12 @@ def test_torch_return_types_returns(self, device):
                 vmap(torch.topk, (0, None, None))(t, 1, 0), torch.return_types.topk
             )
         )
-        self.assertTrue(
-            isinstance(vmap(torch.linalg.eig, (0))(t), torch.return_types.linalg_eig)
-        )
+        if not (TEST_WITH_ROCM and not torch.cuda.has_magma):
+            self.assertTrue(
+                isinstance(
+                    vmap(torch.linalg.eig, (0))(t), torch.return_types.linalg_eig
+                )
+            )
 
     def test_namedtuple_returns(self, device):
         Point = namedtuple("Point", ["x", "y"])