diff --git a/src/tir/ir/data_type_rewriter.cc b/src/tir/ir/data_type_rewriter.cc
index b0457a12398a..346f1ab63250 100644
--- a/src/tir/ir/data_type_rewriter.cc
+++ b/src/tir/ir/data_type_rewriter.cc
@@ -373,18 +373,17 @@ Map<String, ffi::Any> IndexDataTypeRewriter::VisitBlockAnnotations(
     const Map<String, ffi::Any>& annotations) {
   auto new_annotations = annotations;
 
-  std::function<ObjectRef(const ObjectRef&)> f_mutate_obj =
-      [this, &f_mutate_obj](const ObjectRef& obj) -> ObjectRef {
-    if (!obj.defined()) {
+  std::function<Any(const Any&)> f_mutate_obj = [this, &f_mutate_obj](const Any& obj) -> Any {
+    if (obj == nullptr) {
       return obj;
     }
-    if (obj->IsInstance<BufferNode>()) {
+    if (obj.as<BufferNode>()) {
       Buffer buffer = Downcast<Buffer>(obj);
       if (Buffer new_buffer = GetRemappedBuffer(buffer); !new_buffer.same_as(buffer)) {
         return new_buffer;
       }
-    } else if (obj->IsInstance<ffi::ArrayObj>()) {
-      return Downcast<Array<ObjectRef>>(obj).Map(f_mutate_obj);
+    } else if (obj.as<ffi::ArrayObj>()) {
+      return Downcast<Array<Any>>(obj).Map(f_mutate_obj);
     }
     return obj;
   };