fix(mgb): add compute capability check

GitOrigin-RevId: fba819afb3cc154e092f7e9bd84910890d0916ea

fix(mgb): add compute capability check
GitOrigin-RevId: fba819afb3cc154e092f7e9bd84910890d0916ea
ab3efab0 · Megvii Engine Team · Xu Xinran · 8070f40a · ab3efab0 · ab3efab0
隐藏空白更改
内联并排

Showing with 7 addition and 0 deletion

src/gopt/test/inference.cpp src/gopt/test/inference.cpp +3 -0

src/tensorrt/test/opr_replace.cpp src/tensorrt/test/opr_replace.cpp +4 -0

未找到文件。
--- a/src/gopt/test/inference.cpp
+++ b/src/gopt/test/inference.cpp
@@ -2727,6 +2727,9 @@ TEST(TestGoptInference, ConvertFormatNCHW4) {

 TEST(TestGoptInference, ConvertFormatNCHW4Ic3) {
    REQUIRE_GPU(1);
+    auto cn = CompNode::load("gpu0");
+    cn.activate();
+    REQUIRE_CUDA_COMPUTE_CAPABILITY(6, 1);
    HostTensorGenerator<dtype::Float32, RandomDistribution::UNIFORM> gen{
            1.2f, 127 * 127};
    auto graph = ComputingGraph::make();

--- a/src/tensorrt/test/opr_replace.cpp
+++ b/src/tensorrt/test/opr_replace.cpp
@@ -1980,6 +1980,10 @@ TEST(TestTensorRTReplace, FuseConvAdd) {

 TEST(TestTensorRTReplace, FuseConvAddNchw2nchw4) {
    REQUIRE_GPU(1);
+    auto cn = CompNode::load("gpu0");
+    cn.activate();
+    REQUIRE_CUDA_COMPUTE_CAPABILITY(6, 1);
+
    HostTensorGenerator<dtype::Float32, RandomDistribution::UNIFORM> gen{
            1.2f, 127 * 127};
    auto graph = ComputingGraph::make();