fix a bug during folding quantize op (alibaba#968)

Mookel · Jan 19, 2023 · ca7f28e · ca7f28e
1 parent 2e56fa8
commit ca7f28e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
@@ -168,8 +168,10 @@ OpFoldResult QuantizeOp::fold(ArrayRef<Attribute> operands) {
     return {};
   }
 
-  DenseElementsAttr scale = operands[1].dyn_cast<DenseElementsAttr>();
-  DenseElementsAttr zeroPoint = operands[2].dyn_cast<DenseElementsAttr>();
+  DenseElementsAttr scale = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  DenseElementsAttr zeroPoint =
+      operands[2].dyn_cast_or_null<DenseElementsAttr>();
+  if (!scale || !zeroPoint) return {};
   ArrayRef<int64_t> scaleShape = scale.getType().getShape();
   ArrayRef<int64_t> zeroPointShape = zeroPoint.getType().getShape();
   // scale & zero_point must have the same shape