diff --git a/src/transform/layout_inference.cc b/src/transform/layout_inference.cc
index 427549303..628b61ce3 100644
--- a/src/transform/layout_inference.cc
+++ b/src/transform/layout_inference.cc
@@ -159,7 +159,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
           }
         }
         // If already in map, ensure they are structurally equal
-        ICHECK(StructuralEqual()(layout, layout_map[buffer]))
+        ICHECK(layout->IsEqual(layout_map[buffer].get()))
             << "Get different layout for " << buffer
             << "\n current layout: " << layout->DebugOutput()
             << "\n previous layout: " << layout_map[buffer]->DebugOutput();