[AOTI] Not use AOTI_TORCH_CHECK in non AOTI mode. (#143970)

Fix #143967 Pull Request resolved: #143970 Approved by: https://github.com/EikanWang, https://github.com/jansel
pytorch · Dec 31, 2024 · 01034e9 · 01034e9
1 parent a2753e3
commit 01034e9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torch/_inductor/codegen/cpp.py b/torch/_inductor/codegen/cpp.py
@@ -2275,7 +2275,10 @@ def codegen_loops(self, code, worksharing):
 
     @property
     def assert_function(self) -> str:
-        return "AOTI_TORCH_CHECK"
+        if V.graph.aot_mode:
+            return "AOTI_TORCH_CHECK"
+        else:
+            return "TORCH_CHECK"
 
     def decide_parallel_depth(self, max_parallel_depth, threads):
         assert self.call_ranges is not None