Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit 5dffe7a

Browse files
committed
Skip compute capability <= 52 for TVM generated ops
1 parent b49cfce commit 5dffe7a

File tree

1 file changed

+10
-1
lines changed

1 file changed

+10
-1
lines changed

python/mxnet/test_utils.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2232,4 +2232,13 @@ def is_cd_run():
22322232

22332233

22342234
def has_tvm_ops():
2235-
return _features.is_enabled("TVM_OP")
2235+
"""Returns True if MXNet is compiled with TVM generated operators. If current ctx
2236+
is GPU, it only returns True for CUDA compute capability > 52 where FP16 is supported."""
2237+
built_with_tvm_op = _features.is_enabled("TVM_OP")
2238+
if current_context().device_type == 'gpu':
2239+
try:
2240+
import tvm
2241+
except:
2242+
return False
2243+
return built_with_tvm_op and (int("".join(tvm.nd.gpu(0).compute_version.split('.'))) >= 53)
2244+
return built_with_tvm_op

0 commit comments

Comments
 (0)