diff --git a/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc b/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc
index 5cfe9d2e0f6..6c0d626978f 100644
--- a/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc
@@ -53,7 +53,7 @@ LogicalResult ConstantOpConverter::matchAndRewrite(
   auto resultType = op.getType().dyn_cast<RankedTensorType>();
   if (!resultType) return failure();
 
-  if (resultType.getRank() != 1) return failure();
+  if (resultType.getRank() > 1) return failure();
 
   auto elemType = resultType.getElementType();
   if (!elemType.isIndex() && !elemType.isa<IntegerType>()) return failure();
@@ -68,7 +68,12 @@ LogicalResult ConstantOpConverter::matchAndRewrite(
         rewriter.create<arith::ConstantIndexOp>(loc, en.value().getSExtValue());
     if (!elemType.isIndex())
       val = rewriter.create<arith::IndexCastOp>(loc, elemType, val);
-    rewriter.create<memref::StoreOp>(loc, val, result, idx);
+    if (resultType.getRank() == 0) {
+      rewriter.create<memref::StoreOp>(loc, val, result);
+    } else {
+      Value idx = rewriter.create<arith::ConstantIndexOp>(loc, en.index());
+      rewriter.create<memref::StoreOp>(loc, val, result, idx);
+    }
   }
 
   rewriter.replaceOp(op, {result});