[Relay][Quantization] Fix issue introduced in apache#3135

vinx13 · Jul 8, 2019 · 138b0d5 · 138b0d5
1 parent 2a7aebe
commit 138b0d5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/python/tvm/relay/quantize/_annotate.py b/python/tvm/relay/quantize/_annotate.py
@@ -260,7 +260,6 @@ def add_rewrite(ref_call, new_args, ctx):
         if isinstance(rhs_expr, _expr.Constant):
             # quantize rhs to WEIGHT field if it is Constant
             rhs_expr = attach_simulated_quantize(rhs_expr, QAnnotateKind.WEIGHT)
-            assert lhs_kind == QAnnotateKind.ACTIVATION
             expr = _forward_op(ref_call, [lhs_expr, rhs_expr])
             return QAnnotateExpr(expr, QAnnotateKind.ACTIVATION)
         else:

diff --git a/src/relay/pass/quantize.cc b/src/relay/pass/quantize.cc
@@ -413,14 +413,15 @@ Array<Expr> UnifyDTypeScale(const Array<Expr>& ref_args,
   // unify the data type
   CHECK_EQ(ref_args.size(), args.size());
   DataType dtype;
-  if (nptrs[0]->dtype == cfg->dtype_activation) {
-    DataType dtype = cfg->dtype_activation;
-    ret.Set(1, Cast(ret[1], dtype));
-  } else if (nptrs[1]->dtype == cfg->dtype_input) {
-    DataType dtype = cfg->dtype_input;
-    ret.Set(0, Cast(ret[0], dtype));
+  if (ret.size() == 2 && nptrs[1]->dtype == cfg->dtype_input) {
+    dtype = cfg->dtype_input;
   } else {
-    LOG(FATAL) << "should not touch here.";
+    dtype = cfg->dtype_activation;
+  }
+  for (size_t i = 0; i < ret.size(); ++i) {
+    if (nptrs[i]->dtype != dtype) {
+      ret.Set(i, Cast(ret[i], dtype));
+    }
   }
 
   // unify the dom_scale