diff --git a/op-finder/OperationFinderAstVisitor.cpp b/op-finder/OperationFinderAstVisitor.cpp
index aa0b8e7..8d7249d 100644
--- a/op-finder/OperationFinderAstVisitor.cpp
+++ b/op-finder/OperationFinderAstVisitor.cpp
@@ -58,10 +58,6 @@ bool OperationFinderAstVisitor::dataTraverseStmtPre(clang::Stmt* stmt)
       _op_finder->forLoopEntered();
     }
   }
-  else
-  {
-    _checkIfForLoopInitDone(stmt);
-  }
 
   return true;
 }
@@ -70,34 +66,15 @@ bool OperationFinderAstVisitor::dataTraverseStmtPost(clang::Stmt* stmt)
 {
   assert(_context);
 
-  if (clang::dyn_cast<clang::ForStmt>(stmt))
+  if (_loop_init && _loop_init == stmt)
+  {
+    _op_finder->forLoopEntered();
+    _loop_init = nullptr;
+  }
+  else if (clang::dyn_cast<clang::ForStmt>(stmt))
+  {
     _op_finder->forLoopExited();
+  }
 
   return true;
 }
-
-void OperationFinderAstVisitor::_checkIfForLoopInitDone(clang::Stmt* stmt)
-{
-  if (!_loop_init || stmt == _loop_init)
-    return;
-
-  const auto [file, line, column] = resolveLocationsWithLoc(stmt->getEndLoc(), _context->getSourceManager());
-  const auto [for_file, for_line, for_column] = resolveLocationsWithLoc(_loop_init->getEndLoc(), _context->getSourceManager());
-
-  bool is_done = false;
-
-  if (line > for_line)
-  {
-    is_done = true;
-  }
-  else if (line == for_line && column > for_column)
-  {
-    is_done = true;
-  }
-
-  if (is_done)
-  {
-    _loop_init = nullptr;
-    _op_finder->forLoopEntered();
-  }
-}
diff --git a/op-finder/OperationFinderAstVisitor.hpp b/op-finder/OperationFinderAstVisitor.hpp
index cf78f1d..041a927 100644
--- a/op-finder/OperationFinderAstVisitor.hpp
+++ b/op-finder/OperationFinderAstVisitor.hpp
@@ -25,8 +25,6 @@ public:
   bool dataTraverseStmtPost(clang::Stmt* stmt);
 
 private:
-  void _checkIfForLoopInitDone(clang::Stmt* stmt);
-
   clang::ASTContext* _context;
   OperationFinder* _op_finder;