Function test_tensorrt_avg_pool2d

tests/python/relax/test_codegen_tensorrt.py:170–182

()

def test_tensorrt_avg_pool2d():
    @tvm.script.ir_module
    class AvgPool:
        @R.function
        def main(data: R.Tensor((2, 8, 16, 16), "float32")):
            with R.dataflow():
                out = relax.op.nn.avg_pool2d(data, pool_size=(2, 2), strides=(2, 2))
                R.output(out)
            return out

    data = np.random.randn(2, 8, 16, 16).astype("float32")
    patterns = [("tensorrt.nn.avg_pool2d", is_op("relax.nn.avg_pool2d")(wildcard()))]
    _offload_and_compare(AvgPool, {}, patterns, data)


def test_tensorrt_softmax():