diff --git a/python/tvm/relay/op/contrib/dnnl.py b/python/tvm/relay/op/contrib/dnnl.py index 412393eb291a..71ef430ec9c6 100644 --- a/python/tvm/relay/op/contrib/dnnl.py +++ b/python/tvm/relay/op/contrib/dnnl.py @@ -33,7 +33,7 @@ check the attributes of the op and decide if it should be offloaded to DNNL. """ from ... import expr as _expr -from ... import op as reg +from ... import op as _op from .register import register_pattern_table @@ -51,7 +51,7 @@ def _register_external_op_helper(op_name, supported=True): f : callable A function that returns if the operator is supported by DNNL. """ - @reg.register(op_name, "target.dnnl") + @_op.register(op_name, "target.dnnl") def _func_wrapper(attrs, args): return supported @@ -71,12 +71,12 @@ def make_pattern(with_bias=True): data = _expr.var("data") weight = _expr.var("weight") bias = _expr.var("bias") - conv = reg.nn.conv2d(data, weight) + conv = _op.nn.conv2d(data, weight) if with_bias: - conv_out = reg.add(conv, bias) + conv_out = _op.add(conv, bias) else: conv_out = conv - return reg.nn.relu(conv_out) + return _op.nn.relu(conv_out) @register_pattern_table("dnnl")