CR for #24 (#30)

ezyang · May 1, 2024 · 6401e99 · 6401e99
1 parent c79c038
commit 6401e99
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());