Fix flaky nuclear_norm() test (#21638)

Summary: Try to fix a sporadic failure on some CIs. I've run this test hundreds of times on my machine (GeForce 1060, MAGMA) but I cannot reproduce this. Pull Request resolved: https://github.com/pytorch/pytorch/pull/21638 Differential Revision: D15827779 Pulled By: ezyang fbshipit-source-id: 3586075e48907b3b84a101c560a34cc733514a02

Fix flaky nuclear_norm() test (#21638)
Summary: Try to fix a sporadic failure on some CIs. I've run this test hundreds of times on my machine (GeForce 1060, MAGMA) but I cannot reproduce this. Pull Request resolved: https://github.com/pytorch/pytorch/pull/21638 Differential Revision: D15827779 Pulled By: ezyang fbshipit-source-id: 3586075e48907b3b84a101c560a34cc733514a02
71082187 · Stefan Krah · Facebook Github Bot · ff8c3fd5 · 71082187 · 71082187
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -2692,10 +2692,12 @@ class TestCuda(TestCase):

    @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
    @unittest.skipIf(not TEST_MAGMA, "no MAGMA library detected")
+    @skipCUDANonDefaultStreamIf(True)
    def test_nuclear_norm_axes_small_brute_force(self):
        _TestTorchMixin._test_nuclear_norm_axes(self, device='cuda')

    @unittest.skipIf(not TEST_MAGMA, "no MAGMA library detected")
+    @skipCUDANonDefaultStreamIf(True)
    def test_nuclear_norm_exceptions(self):
        _TestTorchMixin._test_nuclear_norm_exceptions(self, device='cuda')


--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -991,14 +991,14 @@ class _TestTorchMixin(object):
            ans = torch.norm(x, "nuc", dim=axes)
            self.assertTrue(ans.is_contiguous())
            self.assertEqual(ans.shape, expected.shape)
-            self.assertTrue(np.allclose(ans.cpu(), expected, rtol=1e-02, atol=1e-03))
+            self.assertTrue(np.allclose(ans.cpu(), expected, rtol=1e-02, atol=1e-03, equal_nan=True))

            out = torch.zeros(expected.shape, dtype=x.dtype, device=x.device)
            ans = torch.norm(x, "nuc", dim=axes, out=out)
            self.assertIs(ans, out)
            self.assertTrue(ans.is_contiguous())
            self.assertEqual(ans.shape, expected.shape)
-            self.assertTrue(np.allclose(ans.cpu(), expected, rtol=1e-02, atol=1e-03))
+            self.assertTrue(np.allclose(ans.cpu(), expected, rtol=1e-02, atol=1e-03, equal_nan=True))

        for n in range(1, 3):
            for m in range(1, 3):
@@ -1026,7 +1026,7 @@ class _TestTorchMixin(object):
                        check_single_nuclear_norm(x, axes)

                        # 3d, inner dimensions Fortran
-                        y = torch.randn(o, n, m, device=device).transpose(-1, -2)
+                        x = torch.randn(o, m, n, device=device).transpose(-1, -2)
                        check_single_nuclear_norm(x, axes)

                        # 3d, inner dimensions non-contiguous