diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc
index 45b8c524650..397889415f7 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc
@@ -159,14 +159,10 @@ class TensorCastOpConverter
   LogicalResult matchAndRewrite(
       TensorCastOp op, ArrayRef<Value> operands,
       ConversionPatternRewriter &rewriter) const final {
-    auto tensor_ty = op.getType().dyn_cast<RankedTensorType>();
-    if (!tensor_ty) return failure();
-
     Value arg = operands.front();
-    auto arg_ty = arg.getType().dyn_cast<MemRefType>();
-    if (!arg_ty) return failure();
+    if (!arg.getType().isa<MemRefType>()) return failure();
 
-    auto result_ty = converter->convertType(tensor_ty);
+    auto result_ty = converter->convertType(op.getType());
     rewriter.replaceOpWithNewOp<MemRefCastOp>(op, arg, result_ty);
 
     return success();