CR for #24

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 200a5aa7b8d6cdd24d76cb2e168df15e309c0c1a ghstack-comment-id: 2088607801 Pull Request resolved: #30
ezyang · May 1, 2024 · 6c6eff2 · 6c6eff2
1 parent c79c038
commit 6c6eff2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -36,7 +36,7 @@ impl Default for ParseConfig {
     }
 }
 
-fn maybe_remove_suffix(frames: &mut Vec<FrameSummary>) {
+fn maybe_remove_convert_frame_suffixes(frames: &mut Vec<FrameSummary>) {
     let target_frames = [
         ("torch/_dynamo/convert_frame.py", "catch_errors"),
         ("torch/_dynamo/convert_frame.py", "_convert_frame"),
@@ -54,7 +54,7 @@ fn maybe_remove_suffix(frames: &mut Vec<FrameSummary>) {
                     && frame.name == target.1
             })
         {
-            frames.truncate(len - 3);
+            frames.truncate(len - target_frames.len());
         }
     }
 }
@@ -313,7 +313,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
 
         if let Some(m) = e.dynamo_start {
             if let Some(mut stack) = m.stack {
-                maybe_remove_suffix(&mut stack);
+                maybe_remove_convert_frame_suffixes(&mut stack);
                 stack_index
                     .borrow_mut()
                     .insert(e.compile_id.clone(), stack.clone());