From 3483135da7eae8dc35b7913af902e49967c6cfb6 Mon Sep 17 00:00:00 2001 From: li_jinxin Date: Thu, 7 Dec 2023 16:37:20 +0800 Subject: [PATCH] add skip if unsupport device. --- torch_npu/testing/common_distributed.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/torch_npu/testing/common_distributed.py b/torch_npu/testing/common_distributed.py index 3a63676225..cf4e4dcb40 100644 --- a/torch_npu/testing/common_distributed.py +++ b/torch_npu/testing/common_distributed.py @@ -38,6 +38,17 @@ def skipIfUnsupportMultiNPU(npu_number_needed): return skip_dec +def skipOpOnlySupport910B(op_name): + def skip_dec(func): + def wrapper(self): + device_name = torch_npu.npu.get_device_name(0)[:10] + if device_name != 'Ascend910B': + raise unittest.SkipTest(f"The op {op_name} is only suppport on 910B, skip this ut for this device type!") + return func(self) + return wrapper + return skip_dec + + def with_comms(func): if func is None: raise RuntimeError("Test function is None.") -- Gitee