fix rocm build

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15945

Differential Revision: D13630505

Pulled By: zdevito

fbshipit-source-id: a4d2ae1370ab475fc1711027c0c9d2a9192be195

diff --git a/test/test_jit.py b/test/test_jit.py
index 8f2cbe0..3cce42f 100644 (file)
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -10416,6 +10416,7 @@ class TestFuser(JitTestCase):
 
     @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows")
     @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA")
+    @skipIfRocm
     def test_fuse_batch_norm(self):
 
         class ResLike(torch.jit.ScriptModule):
@@ -10456,6 +10457,7 @@ class TestFuser(JitTestCase):
 
     @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows")
     @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA")
+    @skipIfRocm
     def test_threshold(self):
         def f(x):
             return torch.threshold(x, 0, -10) + x + x + x