Skip tensorcore test

apache · May 7, 2020 · bbe428c · bbe428c
1 parent 30b4b13
commit bbe428c
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/topi/tests/python/test_topi_conv2d_nhwc_winograd.py b/topi/tests/python/test_topi_conv2d_nhwc_winograd.py
@@ -137,16 +137,16 @@ def test_conv2d_nhwc_winograd_direct():
 
 def test_conv2d_nhwc_winograd_tensorcore():
     """Test the conv2d with winograd for nhwc layout"""
-    print("test_winograd_tensorcore...")
+    if not nvcc.have_tensorcore(tvm.gpu(0).compute_version):
+        return
     verify_conv2d_nhwc(8,  64, 56,  64, 3, 1, 1, bgemm="tensorcore")
     verify_conv2d_nhwc(8, 128, 28, 128, 3, 1, 1, bgemm="tensorcore")
     verify_conv2d_nhwc(8, 256, 14, 256, 3, 1, 1, bgemm="tensorcore")
 
     verify_conv2d_nhwc(2,  64, 56,  64, 3, 1, (1, 1), add_relu=True, bgemm="tensorcore")
     verify_conv2d_nhwc(2,  64, 56,  64, 3, 1, "SAME", add_relu=True, bgemm="tensorcore")
 
+
 if __name__ == "__main__":
     test_conv2d_nhwc_winograd_direct()
-
-    if nvcc.have_tensorcore(tvm.gpu(0).compute_version):
-        test_conv2d_nhwc_winograd_tensorcore()
+    test_conv2d_nhwc_winograd_tensorcore()