make sure we're not counting non-template, non-externs

exclamaforte · exclamaforte · commit 76ffe99c23d4 · 2025-05-07T13:48:33.000-07:00
diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
@@ -1333,7 +1333,16 @@ def fuse(
     @cache_on_self
     def estimate_flops(self) -> int | None:
         # don't increment counters in fused methods so we don't double count
-        fps = list(filter(None, (node.estimate_flops() for node in self.get_nodes())))
+        fps = list(
+            filter(
+                None,
+                (
+                    node.estimate_flops()
+                    for node in self.get_nodes()
+                    if node.is_template() or node.is_extern()
+                ),
+            )
+        )
         if len(fps) == 0:
             return None
         ret = sum(fps)
@@ -1888,7 +1897,16 @@ def get_outputs(self) -> list[SchedulerBuffer]:
     @cache_on_self
     def estimate_flops(self) -> int | None:
         # don't increment counters in fused methods so we don't double count
-        fps = list(filter(None, (node.estimate_flops() for node in self.get_nodes())))
+        fps = list(
+            filter(
+                None,
+                (
+                    node.estimate_flops()
+                    for node in self.get_nodes()
+                    if node.is_template() or node.is_extern()
+                ),
+            )
+        )
         if len(fps) == 0:
             return None
         ret = sum(fps)
