code clean

octoml · Mar 30, 2023 · 9bf08bd · 9bf08bd
1 parent 52864b1
commit 9bf08bd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -202,7 +202,7 @@ def te_silu(x: te.Tensor):
 
 
 @register_legalize("relax.nn.leaky_relu")
-def _nn_softmax(bb: BlockBuilder, call: Call) -> Expr:
+def _nn_leaky_relu(bb: BlockBuilder, call: Call) -> Expr:
     return bb.call_te(topi.nn.leaky_relu, call.args[0], call.attrs.alpha)
 
 

diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
@@ -48,13 +48,11 @@ TVM_REGISTER_GLOBAL("relax.op.nn.leaky_relu").set_body_typed(leaky_relu);
 
 StructInfo InferStructInfoLeakyRelu(const Call& call, const BlockBuilder& ctx) {
   TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
-  if (data_sinfo->IsUnknownNdim()) {
-    return data_sinfo;
-  }
+
   if (!data_sinfo->IsUnknownDtype() &&
       !(data_sinfo->dtype.is_float() || data_sinfo->dtype.is_bfloat16())) {
-    ctx->ReportFatal(Diagnostic::Error(call) <<
-      "LeakyRelu requires the input tensor to have float dtype. However, the given input dtype is "
+    ctx->ReportFatal(Diagnostic::Error(call) << "LeakyRelu requires the input tensor to have float "
+                                                "dtype. However, the given input dtype is "
                                              << data_sinfo->dtype);
   }